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.

Compresión de modelo escalable

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar notebook

Descripción general

En estas anotaciones se muestra cómo comprimir un modelo con TensorFlow Compression.

En el ejemplo que se encuentra abajo, comprimimos los pesos de un clasificador MNIST a un tamaño mucho más pequeño que el de la representación de su punto flotante, a la vez, manteniendo la exactitud en la clasificación. Se hace con un proceso de dos pasos basado en la publicación Scalable Model Compression by Entropy Penalized Reparameterization (Compresión de modelo escalable por reparametrización con sanción por entropía):

• Entrenamiento de un modelo "comprimible" con una sanción por entropía explícita durante el entrenamiento, que favorece la capacidad de compresión de los parámetros del modelo. El peso de la sanción $\lambda$ habilita el control continuo de las compensaciones entre el tamaño del modelo comprimido y su exactitud.

• Codificación del modelo comprimible en un modelo comprimido con un esquema de código que coincide con la sanción, es decir, que la sanción predice bien el tamaño del modelo. De este modo, se garantiza que el método no requiera de varias iteraciones de entrenamiento, compresión ni reentrenamiento del modelo para ajustarlo.

Este método está estrictamente centrado en el tamaño del modelo comprimido, no en la complejidad computacional. Se puede combinar con una técnica como la poda de modelos para reducir el tamaño y la complejidad.

Ejemplo de resultados de compresión en varios modelos:

Modelo (conjunto de datos)	Tamaño del modelo	Relación de comp.	Error de comp. principal (descomp.)
LeNet300-100 (MNIST)	8.56 KB	124x	1.9% (1.6%)
LeNet5-Caffe (MNIST)	2.84 KB	606x	1.0% (0.7%)
VGG-16 (CIFAR-10)	101 KB	590x	10.0% (6.6%)
ResNet-20-4 (CIFAR-10)	128 KB	134x	8.8% (5.0%)
ResNet-18 (ImageNet)	1.97 MB	24x	30.0% (30.0%)
ResNet-50 (ImageNet)	5.49 MB	19x	26.0% (25.0%)

Entre las aplicaciones se incluye:

• Retransmisión/implementación de modelos para dispositivos con tecnología edge a gran escala, con ahorro de ancho de banda en tránsito.

• Comunicación del estado del modelo global a los clientes en aprendizaje federado. La arquitectura del modelo (la cantidad de unidades ocultas, etc.) no cambia en cuanto al modelo inicial, y los clientes pueden seguir aprendiendo con el modelo comprimido.

• Inferir sobre clientes con memoria extremadamente limitada. Durante la inferencia, los pesos de cada capa se pueden descomprimir secuencialmente y descartar inmediatamente después de que se calculan las activaciones.

Preparar

Instalar Tensorflow Compression con 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))"

Importar las dependencias de la librería.

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

Definición y entrenamiento de un clasificador MNIST básico

A fin de comprimir con efectividad las capas densas y convolucionales necesitamos definir las clases de capas personalizadas. Son análogas a las capas bajo tf.keras.layers, pero las subclasificaremos más adelante para implementar efectivamente la reparametrización con sanción por entropía (EPR). Para este mismo objetivo, también agregamos un constructor de copia.

Primero, definimos una capa densa estándar:

class CustomDense(tf.keras.layers.Layer):

  def __init__(self, filters, name="dense"):
    super().__init__(name=name)
    self.filters = filters

  @classmethod
  def copy(cls, other, **kwargs):
    """Returns an instantiated and built layer, initialized from `other`."""
    self = cls(filters=other.filters, name=other.name, **kwargs)
    self.build(None, other=other)
    return self

  def build(self, input_shape, other=None):
    """Instantiates weights, optionally initializing them from `other`."""
    if other is None:
      kernel_shape = (input_shape[-1], self.filters)
      kernel = tf.keras.initializers.GlorotUniform()(shape=kernel_shape)
      bias = tf.keras.initializers.Zeros()(shape=(self.filters,))
    else:
      kernel, bias = other.kernel, other.bias
    self.kernel = tf.Variable(
        tf.cast(kernel, self.variable_dtype), name="kernel")
    self.bias = tf.Variable(
        tf.cast(bias, self.variable_dtype), name="bias")
    self.built = True

  def call(self, inputs):
    outputs = tf.linalg.matvec(self.kernel, inputs, transpose_a=True)
    outputs = tf.nn.bias_add(outputs, self.bias)
    return tf.nn.leaky_relu(outputs)

De un modo similar, una capa convolucional 2D:

class CustomConv2D(tf.keras.layers.Layer):

  def __init__(self, filters, kernel_size,
               strides=1, padding="SAME", name="conv2d"):
    super().__init__(name=name)
    self.filters = filters
    self.kernel_size = kernel_size
    self.strides = strides
    self.padding = padding

  @classmethod
  def copy(cls, other, **kwargs):
    """Returns an instantiated and built layer, initialized from `other`."""
    self = cls(filters=other.filters, kernel_size=other.kernel_size,
               strides=other.strides, padding=other.padding, name=other.name,
               **kwargs)
    self.build(None, other=other)
    return self

  def build(self, input_shape, other=None):
    """Instantiates weights, optionally initializing them from `other`."""
    if other is None:
      kernel_shape = 2 * (self.kernel_size,) + (input_shape[-1], self.filters)
      kernel = tf.keras.initializers.GlorotUniform()(shape=kernel_shape)
      bias = tf.keras.initializers.Zeros()(shape=(self.filters,))
    else:
      kernel, bias = other.kernel, other.bias
    self.kernel = tf.Variable(
        tf.cast(kernel, self.variable_dtype), name="kernel")
    self.bias = tf.Variable(
        tf.cast(bias, self.variable_dtype), name="bias")
    self.built = True

  def call(self, inputs):
    outputs = tf.nn.convolution(
        inputs, self.kernel, strides=self.strides, padding=self.padding)
    outputs = tf.nn.bias_add(outputs, self.bias)
    return tf.nn.leaky_relu(outputs)

Antes de continuar con la compresión del modelo, verifiquemos que podamos entrenar bien a un clasificador regular.

Defina la arquitectura del modelo:

classifier = tf.keras.Sequential([
    CustomConv2D(20, 5, strides=2, name="conv_1"),
    CustomConv2D(50, 5, strides=2, name="conv_2"),
    tf.keras.layers.Flatten(),
    CustomDense(500, name="fc_1"),
    CustomDense(10, name="fc_2"),
], name="classifier")

Cargue los datos de entrenamiento:

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label

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

Finalmente, entrene el modelo:

def train_model(model, training_data, validation_data, **kwargs):
  model.compile(
      optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
      # Uncomment this to ease debugging:
      # run_eagerly=True,
  )
  kwargs.setdefault("epochs", 5)
  kwargs.setdefault("verbose", 1)
  log = model.fit(
      training_data.batch(128).prefetch(8),
      validation_data=validation_data.batch(128).cache(),
      validation_freq=1,
      **kwargs,
  )
  return log.history["val_sparse_categorical_accuracy"][-1]

classifier_accuracy = train_model(
    classifier, training_dataset, validation_dataset)

print(f"Accuracy: {classifier_accuracy:0.4f}")

¡Excelente! El modelo se entrenó bien y alcanzó una exactitud de más del 98 % en el conjunto de validación en 5 épocas.

Entrenamiento de un clasificador comprimible

La reparametrización con sanción por entropía (EPR) tiene dos ingredientes principales:

• La aplicación de una sanción a los pesos del modelo durante el entrenamiento que corresponde a su entropía en un modelo probabilístico, que coincide con el esquema de codificación de pesos. A continuación, definimos un Regularizer Keras que implementa esta sanción.

• La reparametrización de pesos; es decir, llevarlos a la representación latente que es más fácil de comprimir (produce una mejor compensación entre la capacidad de compresión y el desempeño del modelo). Para núcleos convolucionales, se ha demostrado que el dominio de Fourier los representa bien. Para otros parámetros, en el siguiente ejemplo simplemente se usa la cuantificación escalar (redondeo) con un tamaño de paso de cuantificación variable.

Primero, defina la sanción.

En el ejemplo siguiente se usa un modelo de código/probabilístico implementado en la clase tfc.PowerLawEntropyModel inspirado por la publicación Optimizing the Communication-Accuracy Trade-off in Federated Learning with Rate-Distortion Theory (Optimización de la compensación de exactitud-comunicación en aprendizaje federal con la teoría de tasa-distorsión). La sanción se define como: $$\log \Bigl(\frac {|x| + \alpha} \alpha\Bigr),$$ donde $x$ es un parámetro del modelo o es la representación latente y $\alpha$ es una pequeña constante para la estabilidad numérica con valores cercanos a 0.

_ = tf.linspace(-5., 5., 501)
plt.plot(_, tfc.PowerLawEntropyModel(0).penalty(_));

La sanción es, efectivamente, una pérdida de regularización (a veces, denominada "pérdida de peso"). El hecho de que sea cóncava con una cúspide en cero favorece la dispersión del peso. El esquema de código aplicado para comprimir los pesos, un código gamma de Elias, produce códigos de longitud de $1 + \lfloor \log_2 |x| \rfloor$ bits para la magnitud del elemento. Es decir, al coincidir con la sanción y al aplicar esta última se minimiza la longitud de código esperada.

class PowerLawRegularizer(tf.keras.regularizers.Regularizer):

  def __init__(self, lmbda):
    super().__init__()
    self.lmbda = lmbda

  def __call__(self, variable):
    em = tfc.PowerLawEntropyModel(coding_rank=variable.shape.rank)
    return self.lmbda * em.penalty(variable)

# Normalizing the weight of the penalty by the number of model parameters is a
# good rule of thumb to produce comparable results across models.
regularizer = PowerLawRegularizer(lmbda=2./classifier.count_params())

En segundo lugar, defina las subclases de CustomDense y CustomConv2D que tienen la siguiente funcionalidad adicional:

• Toman una instancia del regulador que figura arriba y la aplican a los núcleos y sesgos durante el entrenamiento.

• Definen el núcleo y el sesgo como una @property que realiza la cuantificación con gradientes directos cuando se accede a las variables. Esta exactitud refleja el cálculo que se lleva a cabo más adelante en el modelo comprimido.

• Definen las variables log_step adicionales, que representan el logaritmo del tamaño del paso de cuantificación. Mientras más aproximada sea la cuantificación, menor será el tamaño del modelo, pero también será menor la exactitud. Los tamaños de los pasos de cuantificación se pueden entrenar para cada parámetro del modelo, de modo que al realizar la optimización de la función de pérdida con sanción se determinará cuál es el mejor tamaño para el paso de cuantificación.

El paso de cuantificación se define de la siguiente manera:

def quantize(latent, log_step):
  step = tf.exp(log_step)
  return tfc.round_st(latent / step) * step

Con esto, podemos definir la capa densa:

class CompressibleDense(CustomDense):

  def __init__(self, regularizer, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.regularizer = regularizer

  def build(self, input_shape, other=None):
    """Instantiates weights, optionally initializing them from `other`."""
    super().build(input_shape, other=other)
    if other is not None and hasattr(other, "kernel_log_step"):
      kernel_log_step = other.kernel_log_step
      bias_log_step = other.bias_log_step
    else:
      kernel_log_step = bias_log_step = -4.
    self.kernel_log_step = tf.Variable(
        tf.cast(kernel_log_step, self.variable_dtype), name="kernel_log_step")
    self.bias_log_step = tf.Variable(
        tf.cast(bias_log_step, self.variable_dtype), name="bias_log_step")
    self.add_loss(lambda: self.regularizer(
        self.kernel_latent / tf.exp(self.kernel_log_step)))
    self.add_loss(lambda: self.regularizer(
        self.bias_latent / tf.exp(self.bias_log_step)))

  @property
  def kernel(self):
    return quantize(self.kernel_latent, self.kernel_log_step)

  @kernel.setter
  def kernel(self, kernel):
    self.kernel_latent = tf.Variable(kernel, name="kernel_latent")

  @property
  def bias(self):
    return quantize(self.bias_latent, self.bias_log_step)

  @bias.setter
  def bias(self, bias):
    self.bias_latent = tf.Variable(bias, name="bias_latent")

La capa convolucional es análoga. Además, el núcleo de convolución se almacena como su transformada discreta de Fourier real (RDFT) siempre que el núcleo esté configurado y la transformada esté invertida cada vez que este se use. Debido a que los diferentes componentes de frecuencia del núcleo tienden a ser más o menos comprimibles, cada uno de ellos tiene su propio tamaño de paso de cuantificación asignado.

Defina la transformada de Fourier y su inversa de la siguiente manera:

def to_rdft(kernel, kernel_size):
  # The kernel has shape (H, W, I, O) -> transpose to take DFT over last two
  # dimensions.
  kernel = tf.transpose(kernel, (2, 3, 0, 1))
  # The RDFT has type complex64 and shape (I, O, FH, FW).
  kernel_rdft = tf.signal.rfft2d(kernel)
  # Map real and imaginary parts into regular floats. The result is float32
  # and has shape (I, O, FH, FW, 2).
  kernel_rdft = tf.stack(
      [tf.math.real(kernel_rdft), tf.math.imag(kernel_rdft)], axis=-1)
  # Divide by kernel size to make the DFT orthonormal (length-preserving).
  return kernel_rdft / kernel_size

def from_rdft(kernel_rdft, kernel_size):
  # Undoes the transformations in to_rdft.
  kernel_rdft *= kernel_size
  kernel_rdft = tf.dtypes.complex(*tf.unstack(kernel_rdft, axis=-1))
  kernel = tf.signal.irfft2d(kernel_rdft, fft_length=2 * (kernel_size,))
  return tf.transpose(kernel, (2, 3, 0, 1))

Con esto, defina la capa convolucional como:

class CompressibleConv2D(CustomConv2D):

  def __init__(self, regularizer, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.regularizer = regularizer

  def build(self, input_shape, other=None):
    """Instantiates weights, optionally initializing them from `other`."""
    super().build(input_shape, other=other)
    if other is not None and hasattr(other, "kernel_log_step"):
      kernel_log_step = other.kernel_log_step
      bias_log_step = other.bias_log_step
    else:
      kernel_log_step = tf.fill(self.kernel_latent.shape[2:], -4.)
      bias_log_step = -4.
    self.kernel_log_step = tf.Variable(
        tf.cast(kernel_log_step, self.variable_dtype), name="kernel_log_step")
    self.bias_log_step = tf.Variable(
        tf.cast(bias_log_step, self.variable_dtype), name="bias_log_step")
    self.add_loss(lambda: self.regularizer(
        self.kernel_latent / tf.exp(self.kernel_log_step)))
    self.add_loss(lambda: self.regularizer(
        self.bias_latent / tf.exp(self.bias_log_step)))

  @property
  def kernel(self):
    kernel_rdft = quantize(self.kernel_latent, self.kernel_log_step)
    return from_rdft(kernel_rdft, self.kernel_size)

  @kernel.setter
  def kernel(self, kernel):
    kernel_rdft = to_rdft(kernel, self.kernel_size)
    self.kernel_latent = tf.Variable(kernel_rdft, name="kernel_latent")

  @property
  def bias(self):
    return quantize(self.bias_latent, self.bias_log_step)

  @bias.setter
  def bias(self, bias):
    self.bias_latent = tf.Variable(bias, name="bias_latent")

Defina un modelo clasificador con la misma arquitectura que ve arriba, pero use estas capas modificadas:

def make_mnist_classifier(regularizer):
  return tf.keras.Sequential([
      CompressibleConv2D(regularizer, 20, 5, strides=2, name="conv_1"),
      CompressibleConv2D(regularizer, 50, 5, strides=2, name="conv_2"),
      tf.keras.layers.Flatten(),
      CompressibleDense(regularizer, 500, name="fc_1"),
      CompressibleDense(regularizer, 10, name="fc_2"),
  ], name="classifier")

compressible_classifier = make_mnist_classifier(regularizer)

Y entrene el modelo:

penalized_accuracy = train_model(
    compressible_classifier, training_dataset, validation_dataset)

print(f"Accuracy: {penalized_accuracy:0.4f}")

El modelo comprimible ha alcanzado una exactitud similar a la del clasificador simple.

Sin embargo, el modelo todavía no está realmente comprimido. Para lograrlo, definimos otro conjunto de subclases que almacena los núcleos y los sesgos en forma comprimida, como una secuencia de bits.

Compresión del clasificador

Las subclases de CustomDense y CustomConv2D definidas a continuación convierten los pesos de una capa densa comprimible en cadenas binarias. Además, almacenan el logaritmo del tamaño del paso de la cuantificación a la mitad de la precisión, para ahorrar espacio. Siempre que se accede al núcleo o sesgo mediante @property, se descomprimen a partir de su representación de cadena y se descuantifican.

Primero, defina las funciones para comprimir y descomprimir un parámetro del modelo:

def compress_latent(latent, log_step, name):
  em = tfc.PowerLawEntropyModel(latent.shape.rank)
  compressed = em.compress(latent / tf.exp(log_step))
  compressed = tf.Variable(compressed, name=f"{name}_compressed")
  log_step = tf.cast(log_step, tf.float16)
  log_step = tf.Variable(log_step, name=f"{name}_log_step")
  return compressed, log_step

def decompress_latent(compressed, shape, log_step):
  latent = tfc.PowerLawEntropyModel(len(shape)).decompress(compressed, shape)
  step = tf.exp(tf.cast(log_step, latent.dtype))
  return latent * step

Con esto podemos definir CompressedDense:

class CompressedDense(CustomDense):

  def build(self, input_shape, other=None):
    assert isinstance(other, CompressibleDense)
    self.input_channels = other.kernel.shape[0]
    self.kernel_compressed, self.kernel_log_step = compress_latent(
        other.kernel_latent, other.kernel_log_step, "kernel")
    self.bias_compressed, self.bias_log_step = compress_latent(
        other.bias_latent, other.bias_log_step, "bias")
    self.built = True

  @property
  def kernel(self):
    kernel_shape = (self.input_channels, self.filters)
    return decompress_latent(
        self.kernel_compressed, kernel_shape, self.kernel_log_step)

  @property
  def bias(self):
    bias_shape = (self.filters,)
    return decompress_latent(
        self.bias_compressed, bias_shape, self.bias_log_step)

La clase de capa convolucional es análoga a la anterior.

class CompressedConv2D(CustomConv2D):

  def build(self, input_shape, other=None):
    assert isinstance(other, CompressibleConv2D)
    self.input_channels = other.kernel.shape[2]
    self.kernel_compressed, self.kernel_log_step = compress_latent(
        other.kernel_latent, other.kernel_log_step, "kernel")
    self.bias_compressed, self.bias_log_step = compress_latent(
        other.bias_latent, other.bias_log_step, "bias")
    self.built = True

  @property
  def kernel(self):
    rdft_shape = (self.input_channels, self.filters,
                  self.kernel_size, self.kernel_size // 2 + 1, 2)
    kernel_rdft = decompress_latent(
        self.kernel_compressed, rdft_shape, self.kernel_log_step)
    return from_rdft(kernel_rdft, self.kernel_size)

  @property
  def bias(self):
    bias_shape = (self.filters,)
    return decompress_latent(
        self.bias_compressed, bias_shape, self.bias_log_step)

Para transformar el modelo comprimible en uno comprimido podemos aprovechar y usar la función clone_model. compress_layer convierte cualquier capa comprimible en comprimida y, simplemente, pasa a través de cualquier otro tipo de capa (como Flatten, etc.).

def compress_layer(layer):
  if isinstance(layer, CompressibleDense):
    return CompressedDense.copy(layer)
  if isinstance(layer, CompressibleConv2D):
    return CompressedConv2D.copy(layer)
  return type(layer).from_config(layer.get_config())

compressed_classifier = tf.keras.models.clone_model(
    compressible_classifier, clone_function=compress_layer)

Ahora, validemos que el modelo comprimido aún se comporta como se espera:

compressed_classifier.compile(metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
_, compressed_accuracy = compressed_classifier.evaluate(validation_dataset.batch(128))

print(f"Accuracy of the compressible classifier: {penalized_accuracy:0.4f}")
print(f"Accuracy of the compressed classifier: {compressed_accuracy:0.4f}")

¡La exactitud de clasificación del modelo comprimido es idéntica a la lograda durante el entrenamiento!

Además, el tamaño del modelo comprimido pesa mucho menos que el del tamaño del modelo original:

def get_weight_size_in_bytes(weight):
  if weight.dtype == tf.string:
    return tf.reduce_sum(tf.strings.length(weight, unit="BYTE"))
  else:
    return tf.size(weight) * weight.dtype.size

original_size = sum(map(get_weight_size_in_bytes, classifier.weights))
compressed_size = sum(map(get_weight_size_in_bytes, compressed_classifier.weights))

print(f"Size of original model weights: {original_size} bytes")
print(f"Size of compressed model weights: {compressed_size} bytes")
print(f"Compression ratio: {(original_size/compressed_size):0.0f}x")

Para almacenar los modelos en el disco se requiere de sobrecosto para almacenar la arquitectura del modelo, los gráficos de función del modelo, etc.

Los métodos de compresión sin pérdida como ZIP son buenos para comprimir este tipo de datos, pero no los pesos en sí mismos. Es el motivo por el que sigue siendo beneficioso usar EPR para contar el tamaño del modelo incluido en ese sobrecosto, después de aplicar también la compresión ZIP:

import os
import shutil

def get_disk_size(model, path):
  model.save(path)
  zip_path = shutil.make_archive(path, "zip", path)
  return os.path.getsize(zip_path)

original_zip_size = get_disk_size(classifier, "/tmp/classifier")
compressed_zip_size = get_disk_size(
    compressed_classifier, "/tmp/compressed_classifier")

print(f"Original on-disk size (ZIP compressed): {original_zip_size} bytes")
print(f"Compressed on-disk size (ZIP compressed): {compressed_zip_size} bytes")
print(f"Compression ratio: {(original_zip_size/compressed_zip_size):0.0f}x")

Efecto de la regularización y compensación de tamaño-exactitud

Arriba, el hiperparámetro $\lambda$ se estableció en 2 (normalizado por la cantidad de parámetros del modelo). A medida que aumentamos $\lambda$, los pesos del modelo reciben cada vez más sanciones por la comprimibilidad.

Para valores bajos, la sanción puede funcionar como un regularizador de peso. De hecho, tiene un efecto favorable en el desempeño de la generalización del clasificador y puede contribuir a una exactitud levemente mayor en el conjunto de datos de validación:

#@title

print(f"Accuracy of the vanilla classifier: {classifier_accuracy:0.4f}")
print(f"Accuracy of the penalized classifier: {penalized_accuracy:0.4f}")

Para valores más altos, vemos un tamaño de modelo cada vez más pequeño, pero también la exactitud disminuye gradualmente. Para verlo, entrene algunos modelos y grafique su tamaño con respecto a la exactitud:

def compress_and_evaluate_model(lmbda):
  print(f"lambda={lmbda:0.0f}: training...", flush=True)
  regularizer = PowerLawRegularizer(lmbda=lmbda/classifier.count_params())
  compressible_classifier = make_mnist_classifier(regularizer)
  train_model(
      compressible_classifier, training_dataset, validation_dataset, verbose=0)
  print("compressing...", flush=True)
  compressed_classifier = tf.keras.models.clone_model(
      compressible_classifier, clone_function=compress_layer)
  compressed_size = sum(map(
      get_weight_size_in_bytes, compressed_classifier.weights))
  compressed_zip_size = float(get_disk_size(
      compressed_classifier, "/tmp/compressed_classifier"))
  print("evaluating...", flush=True)
  compressed_classifier = tf.keras.models.load_model(
      "/tmp/compressed_classifier")
  compressed_classifier.compile(
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
  _, compressed_accuracy = compressed_classifier.evaluate(
      validation_dataset.batch(128), verbose=0)
  print()
  return compressed_size, compressed_zip_size, compressed_accuracy

lambdas = (2., 5., 10., 20., 50.)
metrics = [compress_and_evaluate_model(l) for l in lambdas]
metrics = tf.convert_to_tensor(metrics, tf.float32)

#@title

def plot_broken_xaxis(ax, compressed_sizes, original_size, original_accuracy):
  xticks = list(range(
      int(tf.math.floor(min(compressed_sizes) / 5) * 5),
      int(tf.math.ceil(max(compressed_sizes) / 5) * 5) + 1,
      5))
  xticks.append(xticks[-1] + 10)
  ax.set_xlim(xticks[0], xticks[-1] + 2)
  ax.set_xticks(xticks[1:])
  ax.set_xticklabels(xticks[1:-1] + [f"{original_size:0.2f}"])
  ax.plot(xticks[-1], original_accuracy, "o", label="float32")

sizes, zip_sizes, accuracies = tf.transpose(metrics)
sizes /= 1024
zip_sizes /= 1024

fig, (axl, axr) = plt.subplots(1, 2, sharey=True, figsize=(10, 4))
axl.plot(sizes, accuracies, "o-", label="EPR compressed")
axr.plot(zip_sizes, accuracies, "o-", label="EPR compressed")
plot_broken_xaxis(axl, sizes, original_size/1024, classifier_accuracy)
plot_broken_xaxis(axr, zip_sizes, original_zip_size/1024, classifier_accuracy)

axl.set_xlabel("size of model weights [kbytes]")
axr.set_xlabel("ZIP compressed on-disk model size [kbytes]")
axl.set_ylabel("accuracy")
axl.legend(loc="lower right")
axr.legend(loc="lower right")
axl.grid()
axr.grid()
for i in range(len(lambdas)):
  axl.annotate(f"$\lambda = {lambdas[i]:0.0f}$", (sizes[i], accuracies[i]),
               xytext=(10, -5), xycoords="data", textcoords="offset points")
  axr.annotate(f"$\lambda = {lambdas[i]:0.0f}$", (zip_sizes[i], accuracies[i]),
               xytext=(10, -5), xycoords="data", textcoords="offset points")
plt.tight_layout()

Lo ideal sería que en el gráfico se mostrara una compensación entre tamaño y exactitud con forma de codo, pero lo normal es que las métricas de exactitud, en cierto modo, tengan ruido. Dependiendo de la inicialización, la curva puede exhibir algunos pliegues.

Debido al efecto de regularización, el modelo EPR comprimido es más exacto en el conjunto de prueba que en el modelo original para valores pequeños de $\lambda$. El modelo EPR comprimido también es mucho más chico en comparación con los tamaños posteriores a la compresión ZIP adicional.

Descompresión del clasificador

CompressedDense y CompressedConv2D descomprimen los pesos de cada pase hacia adelante. De este modo, se vuelven ideales para dispositivos de memoria limitada, pero la descompresión puede ser costosa desde el punto de vista del cálculo, en particular en tamaños de lotes pequeños.

Para descomprimir el modelo una vez y usarlo para más entrenamientos o inferencias, podemos volver a convertirlo en modelo con capas regulares o comprimibles. Puede resultar útil para la implementación de modelos o en escenarios de entrenamiento federado.

Primero, lo volvemos a convertir en modelo simple, podemos hacer inferencias o continuar el entrenamiento regular sin una sanción por compresión:

def decompress_layer(layer):
  if isinstance(layer, CompressedDense):
    return CustomDense.copy(layer)
  if isinstance(layer, CompressedConv2D):
    return CustomConv2D.copy(layer)
  return type(layer).from_config(layer.get_config())

decompressed_classifier = tf.keras.models.clone_model(
    compressed_classifier, clone_function=decompress_layer)

decompressed_accuracy = train_model(
    decompressed_classifier, training_dataset, validation_dataset, epochs=1)

print(f"Accuracy of the compressed classifier: {compressed_accuracy:0.4f}")
print(f"Accuracy of the decompressed classifier after one more epoch of training: {decompressed_accuracy:0.4f}")

Tenga en cuenta que la exactitud en la validación cae después de entrenar durante una época adicional, ya que el entrenamiento se hace sin regularización.

Como alternativa, podemos volver a convertir el modelo en uno "comprimible", para inferencias o más entrenamiento con una sanción por compresión:

def decompress_layer_with_penalty(layer):
  if isinstance(layer, CompressedDense):
    return CompressibleDense.copy(layer, regularizer=regularizer)
  if isinstance(layer, CompressedConv2D):
    return CompressibleConv2D.copy(layer, regularizer=regularizer)
  return type(layer).from_config(layer.get_config())

decompressed_classifier = tf.keras.models.clone_model(
    compressed_classifier, clone_function=decompress_layer_with_penalty)

decompressed_accuracy = train_model(
    decompressed_classifier, training_dataset, validation_dataset, epochs=1)

print(f"Accuracy of the compressed classifier: {compressed_accuracy:0.4f}")
print(f"Accuracy of the decompressed classifier after one more epoch of training: {decompressed_accuracy:0.4f}")

Aquí, la exactitud mejora después de entrenar durante una época adicional.