data_compression.ipynb

Kernel: Python 3

Copyright 2022 The TensorFlow Compression Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Compactação de dados aprendidos

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

Este notebook mostra como fazer uma compactação de dados com perda usando redes neurais e o TensorFlow Compression.

A compactação com perda envolve equilibrar a taxa – o número esperado de bits necessários para codificar uma amostra – e a distorção – o erro esperado na reconstrução da amostra.

Os exemplos abaixo usam um modelo tipo autoencoder para compactar imagens do dataset MNIST. O método baseia-se no artigo Compactação de imagens de ponta a ponta.

Confira mais detalhes sobre a compactação de dados aprendidos neste artigo, criado para quem já conhece a compactação de dados clássica, ou nesta pesquisa, para um público de aprendizado de máquina.

Configuração

Instale o TensorFlow Compression pelo pip.

%%bash
# Installs the latest version of TFC compatible with the installed TF version.

read MAJOR MINOR <<< "$(pip show tensorflow | perl -p -0777 -e 's/.*Version: (\d+)\.(\d+).*/\1 \2/sg')"
pip install "tensorflow-compression<$MAJOR.$(($MINOR+1))"

Importe as dependências de biblioteca.

import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_compression as tfc
import tensorflow_datasets as tfds

Definir o modelo do treinador

Como o modelo parece um autoencoder e precisamos realizar um conjunto diferente de funções durante o treinamento e a inferência, a configuração é um pouco diferente de um classificador, por exemplo.

O modelo do treinamento é composto por três partes:

• A transformação de análise (ou encoder), convertendo a imagem em um espaço latente.

• A transformação de síntese (ou decoder), convertendo o espaço latente de volta no espaço de imagem.

• Um modelo de prior e entropia, que modela as probabilidades marginais dos latentes.

Primeiro, defina as transformações:

def make_analysis_transform(latent_dims):
  """Creates the analysis (encoder) transform."""
  return tf.keras.Sequential([
      tf.keras.layers.Conv2D(
          20, 5, use_bias=True, strides=2, padding="same",
          activation="leaky_relu", name="conv_1"),
      tf.keras.layers.Conv2D(
          50, 5, use_bias=True, strides=2, padding="same",
          activation="leaky_relu", name="conv_2"),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(
          500, use_bias=True, activation="leaky_relu", name="fc_1"),
      tf.keras.layers.Dense(
          latent_dims, use_bias=True, activation=None, name="fc_2"),
  ], name="analysis_transform")

def make_synthesis_transform():
  """Creates the synthesis (decoder) transform."""
  return tf.keras.Sequential([
      tf.keras.layers.Dense(
          500, use_bias=True, activation="leaky_relu", name="fc_1"),
      tf.keras.layers.Dense(
          2450, use_bias=True, activation="leaky_relu", name="fc_2"),
      tf.keras.layers.Reshape((7, 7, 50)),
      tf.keras.layers.Conv2DTranspose(
          20, 5, use_bias=True, strides=2, padding="same",
          activation="leaky_relu", name="conv_1"),
      tf.keras.layers.Conv2DTranspose(
          1, 5, use_bias=True, strides=2, padding="same",
          activation="leaky_relu", name="conv_2"),
  ], name="synthesis_transform")

O treinador armazena uma instância das duas transformações, bem como os parâmetros do prior.

Seu método call é configurado para computar:

• A taxa, uma estimativa do número de bits necessários para representar o lote de dígitos.

• A distorção, a diferença absoluta média entre os pixels dos dígitos originais e suas reconstruções.

class MNISTCompressionTrainer(tf.keras.Model):
  """Model that trains a compressor/decompressor for MNIST."""

  def __init__(self, latent_dims):
    super().__init__()
    self.analysis_transform = make_analysis_transform(latent_dims)
    self.synthesis_transform = make_synthesis_transform()
    self.prior_log_scales = tf.Variable(tf.zeros((latent_dims,)))

  @property
  def prior(self):
    return tfc.NoisyLogistic(loc=0., scale=tf.exp(self.prior_log_scales))

  def call(self, x, training):
    """Computes rate and distortion losses."""
    # Ensure inputs are floats in the range (0, 1).
    x = tf.cast(x, self.compute_dtype) / 255.
    x = tf.reshape(x, (-1, 28, 28, 1))

    # Compute latent space representation y, perturb it and model its entropy,
    # then compute the reconstructed pixel-level representation x_hat.
    y = self.analysis_transform(x)
    entropy_model = tfc.ContinuousBatchedEntropyModel(
        self.prior, coding_rank=1, compression=False)
    y_tilde, rate = entropy_model(y, training=training)
    x_tilde = self.synthesis_transform(y_tilde)

    # Average number of bits per MNIST digit.
    rate = tf.reduce_mean(rate)

    # Mean absolute difference across pixels.
    distortion = tf.reduce_mean(abs(x - x_tilde))

    return dict(rate=rate, distortion=distortion)

Computar a taxa e a distorção

Mostraremos o passo a passo, usando uma imagem do conjunto de treinamento. Carregue o dataset MNIST para treinamento e validação:

training_dataset, validation_dataset = tfds.load(
    "mnist",
    split=["train", "test"],
    shuffle_files=True,
    as_supervised=True,
    with_info=False,
)

Extraia uma imagem $x$:

(x, _), = validation_dataset.take(1)

plt.imshow(tf.squeeze(x))
print(f"Data type: {x.dtype}")
print(f"Shape: {x.shape}")

Para obter a representação latente $y$, precisamos convertê-la em float32, adicionar uma dimensão de lote e passá-la pela transformação de análise.

x = tf.cast(x, tf.float32) / 255.
x = tf.reshape(x, (-1, 28, 28, 1))
y = make_analysis_transform(10)(x)

print("y:", y)

Os latentes serão quantizados no momento do teste. Para modelar de uma maneira diferenciável durante o treinamento, adicionamos ruído uniforme no intervalo $(-.5, .5)$ e chamamos o resultado de $\tilde y$. Essa é a mesma terminologia usada no artigo Compactação de imagens de ponta a ponta.

y_tilde = y + tf.random.uniform(y.shape, -.5, .5)

print("y_tilde:", y_tilde)

O “prior” é uma densidade de probabilidade que treinamos para modelar a distribuição marginal dos latentes com ruído. Por exemplo, poderia ser um conjunto de distribuições logísticas independentes com escalas diferentes para cada dimensão de latente. tfc.NoisyLogistic leva em consideração que os latentes têm ruído aditivo. À medida que a escala se aproxima de zero, uma distribuição logística se aproxima de um delta de Dirac (pico), mas o ruído adicionado faz a distribuição “com ruído” se aproximar de uma distribuição uniforme.

prior = tfc.NoisyLogistic(loc=0., scale=tf.linspace(.01, 2., 10))

_ = tf.linspace(-6., 6., 501)[:, None]
plt.plot(_, prior.prob(_));

Durante o treinamento, tfc.ContinuousBatchedEntropyModel adiciona ruído uniforme e usa o ruído e o prior para computar um limite superior (diferenciável) da taxa (o número médio de bits necessários para codificar a representação latente). Esse limite pode ser minimizado como uma perda.

entropy_model = tfc.ContinuousBatchedEntropyModel(
    prior, coding_rank=1, compression=False)
y_tilde, rate = entropy_model(y, training=True)

print("rate:", rate)
print("y_tilde:", y_tilde)

Por fim, os latentes com ruído são passados de volta pela transformação de síntese para produzir uma reconstrução da imagem $\tilde x$. A distorção é o erro entre a imagem original e a reconstrução. Obviamente, com as transformações não treinadas, a reconstrução não é muito útil.

x_tilde = make_synthesis_transform()(y_tilde)

# Mean absolute difference across pixels.
distortion = tf.reduce_mean(abs(x - x_tilde))
print("distortion:", distortion)

x_tilde = tf.saturate_cast(x_tilde[0] * 255, tf.uint8)
plt.imshow(tf.squeeze(x_tilde))
print(f"Data type: {x_tilde.dtype}")
print(f"Shape: {x_tilde.shape}")

Para cada lote de dígitos, uma chamada a MNISTCompressionTrainer produz a taxa e a distorção como uma média para esse lote:

(example_batch, _), = validation_dataset.batch(32).take(1)
trainer = MNISTCompressionTrainer(10)
example_output = trainer(example_batch)

print("rate: ", example_output["rate"])
print("distortion: ", example_output["distortion"])

Na próxima seção, vamos configurar o modelo para fazer o método do gradiente descendente nessas duas perdas.

Treinar o modelo

Compilamos o treinador de uma forma que ele otimiza a taxa-distorção de Lagrange, ou seja, uma soma da taxa e da distorção, em que um dos termos é ponderado pelo parâmetro de Lagrange $\lambda$.

Essa função de perda afeta partes diferentes do modelo de forma diferente:

• A transformação de análise é treinada para produzir uma representação latente que atinge o equilíbrio desejado entre taxa e distorção.

• A transformação de síntese é treinada para minimizar a distorção, dada a representação latente.

• Os parâmetros do prior são treinados para minimizar a taxa, dada a representação latente. Isso é idêntico a colocar o prior na distribuição marginal de latentes quanto à verossimilhança máxima.

def pass_through_loss(_, x):
  # Since rate and distortion are unsupervised, the loss doesn't need a target.
  return x

def make_mnist_compression_trainer(lmbda, latent_dims=50):
  trainer = MNISTCompressionTrainer(latent_dims)
  trainer.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
    # Just pass through rate and distortion as losses/metrics.
    loss=dict(rate=pass_through_loss, distortion=pass_through_loss),
    metrics=dict(rate=pass_through_loss, distortion=pass_through_loss),
    loss_weights=dict(rate=1., distortion=lmbda),
  )
  return trainer

Agora, treine o modelo. Anotações humanas não são necessárias aqui, já que queremos compactar as imagens, então as colocamos em um map e adicionamos alvos “simulados” para a taxa e a distorção.

def add_rd_targets(image, label):
  # Training is unsupervised, so labels aren't necessary here. However, we
  # need to add "dummy" targets for rate and distortion.
  return image, dict(rate=0., distortion=0.)

def train_mnist_model(lmbda):
  trainer = make_mnist_compression_trainer(lmbda)
  trainer.fit(
      training_dataset.map(add_rd_targets).batch(128).prefetch(8),
      epochs=15,
      validation_data=validation_dataset.map(add_rd_targets).batch(128).cache(),
      validation_freq=1,
      verbose=1,
  )
  return trainer

trainer = train_mnist_model(lmbda=2000)

Compactar algumas imagens do MNIST

Para a compactação e descompactação no momento do teste, dividimos o modelo treinado em duas partes:

• O lado do encoder é composto pela transformação de análise e pelo modelo de entropia.

• O lado do decoder é composto pela transformação de síntese e pelo mesmo modelo de entropia.

No momento do teste, os latentes não terão ruído aditivo, mas serão quantizados e depois compactados sem perda, então damos novos nomes a eles. Fazemos uma chamada a eles e à reconstrução da imagem $\hat x$ e $\hat y$, respectivamente (conforme o artigo Compactação de imagens de ponta a ponta).

class MNISTCompressor(tf.keras.Model):
  """Compresses MNIST images to strings."""

  def __init__(self, analysis_transform, entropy_model):
    super().__init__()
    self.analysis_transform = analysis_transform
    self.entropy_model = entropy_model

  def call(self, x):
    # Ensure inputs are floats in the range (0, 1).
    x = tf.cast(x, self.compute_dtype) / 255.
    y = self.analysis_transform(x)
    # Also return the exact information content of each digit.
    _, bits = self.entropy_model(y, training=False)
    return self.entropy_model.compress(y), bits

class MNISTDecompressor(tf.keras.Model):
  """Decompresses MNIST images from strings."""

  def __init__(self, entropy_model, synthesis_transform):
    super().__init__()
    self.entropy_model = entropy_model
    self.synthesis_transform = synthesis_transform

  def call(self, string):
    y_hat = self.entropy_model.decompress(string, ())
    x_hat = self.synthesis_transform(y_hat)
    # Scale and cast back to 8-bit integer.
    return tf.saturate_cast(tf.round(x_hat * 255.), tf.uint8)

Quando instanciado com compression=True, o modelo de entropia converte o prior aprendido em tabelas, para um algoritmo de codificação de intervalo. Quando fazemos uma chamada a compress(), esse algoritmo é chamado para converter o vetor de espaço latente em sequências de bits. O tamanho de cada string binária aproxima o conteúdo de informação do latente (a verossimilhança logarítmica negativa do latente para o prior).

O modelo de entropia para compactação e descompactação deve ser a mesma instância, pois as tabelas de codificação de intervalo precisam ser idênticas nos dois lados. Caso contrário, poderão ocorrer erros de decodificação.

def make_mnist_codec(trainer, **kwargs):
  # The entropy model must be created with `compression=True` and the same
  # instance must be shared between compressor and decompressor.
  entropy_model = tfc.ContinuousBatchedEntropyModel(
      trainer.prior, coding_rank=1, compression=True, **kwargs)
  compressor = MNISTCompressor(trainer.analysis_transform, entropy_model)
  decompressor = MNISTDecompressor(entropy_model, trainer.synthesis_transform)
  return compressor, decompressor

compressor, decompressor = make_mnist_codec(trainer)

Pegue 16 imagens do dataset de validação. Você pode selecionar um subconjunto diferente alterando o argumento para skip.

(originals, _), = validation_dataset.batch(16).skip(3).take(1)

Compacte-as em strings e controle o conteúdo de informação de cada uma delas em bits.

strings, entropies = compressor(originals)

print(f"String representation of first digit in hexadecimal: 0x{strings[0].numpy().hex()}")
print(f"Number of bits actually needed to represent it: {entropies[0]:0.2f}")

Descompacte as imagens das strings.

reconstructions = decompressor(strings)

Exiba cada um dos 16 dígitos originais junto com sua representação binária compactada e o dígito reconstruído.

#@title

def display_digits(originals, strings, entropies, reconstructions):
  """Visualizes 16 digits together with their reconstructions."""
  fig, axes = plt.subplots(4, 4, sharex=True, sharey=True, figsize=(12.5, 5))
  axes = axes.ravel()
  for i in range(len(axes)):
    image = tf.concat([
        tf.squeeze(originals[i]),
        tf.zeros((28, 14), tf.uint8),
        tf.squeeze(reconstructions[i]),
    ], 1)
    axes[i].imshow(image)
    axes[i].text(
        .5, .5, f"→ 0x{strings[i].numpy().hex()} →\n{entropies[i]:0.2f} bits",
        ha="center", va="top", color="white", fontsize="small",
        transform=axes[i].transAxes)
    axes[i].axis("off")
  plt.subplots_adjust(wspace=0, hspace=0, left=0, right=1, bottom=0, top=1)

display_digits(originals, strings, entropies, reconstructions)

O tamanho da string codificada difere do conteúdo de informação de cada dígito.

Isso ocorre porque o processo de codificação de intervalo funciona com probabilidades discretas e uma pequena quantidade de sobrecarga. Portanto, especificamente para strings pequenas, a correspondência é apenas aproximada. Entretanto, a codificação de intervalo é ideal de forma assintótica: no limite, a quantidade esperada de bits se aproximará da entropia cruzada (o conteúdo de informação esperado), em que a taxa no modelo de treinamento é um limite superior.

Equilíbrio entre taxa e distorção

Acima, o modelo foi treinado para um equilíbrio específico (dado lmbda=2000) entre o número médio de bits usados para representar cada dígito e o erro decorrente na reconstrução.

O que acontece quando repetimos o experimento com valores diferentes?

Vamos começar reduzindo $\lambda$ para 500.

def train_and_visualize_model(lmbda):
  trainer = train_mnist_model(lmbda=lmbda)
  compressor, decompressor = make_mnist_codec(trainer)
  strings, entropies = compressor(originals)
  reconstructions = decompressor(strings)
  display_digits(originals, strings, entropies, reconstructions)

train_and_visualize_model(lmbda=500)

A taxa de bits do nosso código cai, assim como a fidelidade dos dígitos. Entretanto, a maioria dos dígitos permanece reconhecível.

Vamos reduzir $\lambda$ ainda mais.

train_and_visualize_model(lmbda=300)

As strings começam a ficar muito menores agora, da ordem de um byte por dígito. Entretanto, isso acarreta um custo. Mais dígitos estão ficando irreconhecíveis.

Isso demonstra que esse modelo é agnóstico quanto às percepções humanas de erro; ele apenas mensura o desvio absoluto quando aos valores de pixels. Para conseguir uma qualidade de imagem melhor, teríamos que substituir a perda de pixel por uma perda perceptiva.

Usar o decoder como modelo generativo

Se alimentarmos bits aleatórios no decoder, será feita uma amostragem da distribuição que o modelo aprendeu para representar dígitos.

Primeiro, instancie novamente o compactador/descompactador sem uma verificação que detectaria se a string de entrada não está totalmente decodificada.

compressor, decompressor = make_mnist_codec(trainer, decode_sanity_check=False)

Agora, alimente strings aleatórias grandes o suficiente no descompactador para que ele possa decodificar/amostrar dígitos.

import os

strings = tf.constant([os.urandom(8) for _ in range(16)])
samples = decompressor(strings)

fig, axes = plt.subplots(4, 4, sharex=True, sharey=True, figsize=(5, 5))
axes = axes.ravel()
for i in range(len(axes)):
  axes[i].imshow(tf.squeeze(samples[i]))
  axes[i].axis("off")
plt.subplots_adjust(wspace=0, hspace=0, left=0, right=1, bottom=0, top=1)