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.

Comprensión de datos aprendida

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver código fuente en GitHub

Descargar bloc de notas

Descripción general

En este cuaderno se muestra como realizar la compresión con pérdida de datos con redes neuronales y TensorFlow Compression.

La compresión con pérdida de datos implica una compensación entre la tasa, el número previsto de bits que se necesitan para codificar una muestra, y la distorsión, el error previsto en la reconstrucción de la muestra.

Los ejemplos siguientes usan un modelo de tipo autocodificador que comprime imágenes del conjunto de datos MNIST. El método se basa en el artículo Compresión optimizada de imágenes de extremo a extremo.

Puede encontrar más información sobre la compresión aprendida de datos en este artículo destinado a personas con conocimientos de la compresión de datos clásica o en este estudio para la audiencia de aprendizaje automático.

Preparación

Instale Tensorflow Compression a través de 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 las dependencias de la biblioteca.

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

Definir el modelo de entrenamiento.

Ya que el modelo se parece a un autocodificador y necesitamos que realice un conjunto de funciones diferentes durante el entrenamiento e inferencia, la instalación es un poco diferente a la de un clasificador, por ejemplo.

El modelo de entrenamiento consiste en tres partes:

• la transformación de análisis (o codificador), que convierte la imagen en un espacio latente,

• la transformación de síntesis (o decodificador), que vuelve a convertir el espacio latente en un espacio de imagen, y

• un modelo de inferencia y entrópico, que modela las probabilidades marginales de los latentes.

Primero, defina las transformaciones:

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")

El entrenador conserva una instancia de ambas transformaciones y también los parámetros de la inferencia.

Se instala su método call para calcular:

• la tasa, un cálculo estimado de la cantidad de bits necesarios para representar un lote de cifras y

• la distorsión, la diferencia media absoluta entre los píxeles de las cifras originales y sus reconstrucciones.

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)

Calcular la tasa y la distorsión.

Vamos paso a paso, usaremos una imagen del conjunto de datos de entrenamiento. Cargue el conjunto de datos MNIST para el entrenamiento y la validación:

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

Y extraiga una imagen $x$:

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

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

Para obtener la representación latente $y$, necesitamos convertirla en float32, agregar una dimensión del lote y pasarla por la transformación de análisis.

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

print("y:", y)

Los latentes serán cuantificados durante el periodo de prueba. Para modelarlo de manera diferenciable durante el entrenamiento, agregamos ruido uniforme en el intervalo $(-.5, .5)$ y llamamos al resultado $\tilde y$. Esta es la misma terminología que se usa en el artículo Compresión optimizada de imágenes de extremo a extremo.

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

print("y_tilde:", y_tilde)

La "inferencia" es la densidad de la probabilidad que entrenamos para modelar la distribución marginal de los latentes con ruido. Por ejemplo, puede ser un conjunto de distribuciones logísticas con diferentes escalas para cada dimensión latente. tfc.NoisyLogistic explica el hecho de que las latentes tengan ruido agregado. Ya que la escala está cerca del cero, una distribución logística se encuentra cerca de la delta de Dirac (pico), pero el ruido agregado causa la distribución "ruidosa" para acercarse más a la distribución uniforme.

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

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

Durante el entrenamiento, tfc.ContinuousBatchedEntropyModel agrega ruido uniforme y usa el ruido y la inferencia para calcular un límite superior (diferenciable) en la tasa (cantidad promedio de bits necesarios para codificar la representación latente). Ese límite puede minimizarse como una pérdida.

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 último, se pasan las latentes con ruido a través de la transformación de síntesis para producir la reconstrucción de la imagen $\tilde x$. La distorsión es el error entre la imagen original y la reconstrucción. Obviamente, si no se entrenan las transformaciones, la reconstrucción no es muy ú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 cifras, si se llama a MNISTCompressionTrainer, este produce la tasa y la distorsión como un promedio sobre ese 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"])

En la siguiente sección, instalaremos el modelo para hacer que el gradiente descienda en las dos pérdidas.

Entrenar el modelo.

Compilamos el entrenador para que optimice la tasa y la distorsión lagareanas, es decir, la suma de la tasa y la distorsión, donde uno de los términos se evalúa con el parámetro $\lambda$ de Lagrange.

Esta función de pérdida afecta a diferentes partes del modelo de distintas formas:

• La transformación de análisis se entrena para producir una representación latente que logre la compensación deseada entre la tasa y la distorsión.

• La transformación de síntesis se entrena para minimizar la distorsión, según la representación latente.

• Los parámetros de la inferencia se entrenan para minimizar la tasa según la representación latente. Es idéntico a encajar la inferencia en la distribución marginal de latentes en un sentido de probabilidad 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

Luego, entrene el modelo. Las anotaciones de personas no son necesarias aquí, ya que solo queremos comprimir las imágenes, por eso las abandonamos con map y en su lugar agregamos destinos "de relleno" para la tasa y la distorsión.

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)

Comprimir algunas imágenes MNIST.

Para comprimir y descomprimir durante el periodo de prueba, dividimos el modelo de entrenamiento en dos partes:

• La parte del codificador que consiste en la transformación de análisis y el modelo de entropía.

• La parte del decodificador que consiste en la transformación de síntesis y el mismo modelo de entropía.

En este momento, las latentes no tienen ruido agregado, pero se cuantificarán y luego se comprimirán sin pérdida así que pondremos nombres nuevos. Los llamamos y a la imagen de reconstrucción $\hat x$ y $\hat y$, respectivamente (como se presenta en Comprensión optimizada de imágenes de extremo a extremo).

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)

Cuando se crea una instancia con compression=True, el modelo de entropía convierte la inferencia aprendida en tablas para un algoritmo de codificación de rango. Cuando se llama a compress(), se invoca este algoritmo para convertir el vector del espacio latente en secuencias de bit. El largo de cada cadena de texto binaria se aproxima al contenido de datos de la latente (la probabilidad logarítmica negativa de la latente bajo la inferencia).

El modelo de entropía para comprimir y descomprimir debe ser la misma instancia, porque las tablas de codificación de rango tienen que ser exactamente idénticas en ambos lados. Si no es así, pueden ocurrir errores de decodificación.

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)

Tome 16 imágenes del conjunto de datos de validación. Puede seleccionar un subconjunto diferente si cambia el argumento a skip.

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

Comprímalas en cadenas de texto y lleve la cuenta de sus contenidos de datos en 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}")

Descomprima las cadena de texto a imágenes.

reconstructions = decompressor(strings)

Muestre cada una de las 16 cifras originales junto a sus representaciones binarias y la cifra reconstruida.

#@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)

Preste atención a como a longitud de la cadena codificada es diferente al contenido de datos de cada cifra.

Esto se debe a que el proceso de codificación de rango funciona con probabilidades discretas y tiene una breve sobrecarga. Por lo tanto, especialmente para las cadenas de texto cortas, la correspondencia solo es aproximada. Sin embargo, la codificación de rango es óptima de forma asintótica: en el límite, el conteo esperado de bits se acercará a la entropía cruzada (el contenido de datos esperado) para el cual el término de la tasa en el modelo de entrenamiento es un límite superior.

La compensación de tasa y distorsión

Anteriormente, se entrenó el modelo para una compensación específica (dada por lmbda=2000) entre la cantidad promedio de bits usados para representar cada cifra y el error resultante durante la reconstrucción.

¿Qué sucede si intentamos el experimento de nuevo con diferentes valores?

Empecemos por reducir $\lambda$ a 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)

La tasa de bits de nuestro código disminuye, al igual que la fidelidad de las cifras. Sin embargo, la mayoría de las cifras siguen siendo reconocibles.

Reduzcamos $\lambda$ un poco más.

train_and_visualize_model(lmbda=300)

Ahora, las cadenas de texto se van acortando, un byte por cifra. Pero esto sucede a un costo. Cada vez más cifras se vuelven irreconocibles.

Esto demuestra que este modelo es agnóstico a la capacidad humana de reconocer los errores, solo mide la desviación absoluta en términos de valores de píxeles. Para lograr una calidad de imagen mejor percibida, tendremos que reemplazar la pérdida de píxeles con la pérdida de percepción.

Usar el decodificador como modelo generativo.

Si ingresamos bits aleatorios en el decodificador, este tomará efectivamente la distribución que el modelo aprendió para representar cifras como muestra.

Primero, vuelva a crear la instancia del compresor/descompresor sin una verificación de sanidad que pueda detectar si la cadena de texto de entrada fue decodificada completamente.

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

Ahora, ingrese cadenas de texto aleatorias con una buena longitud en el descompresor para que pueda decodificar/tomar la muestra de las cifras.

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)