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.

Compressão de modelo escalável

Ver em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Visão geral

Este notebook mostra como fazer a compressão de um modelo usando o TensorFlow Compression.

No exemplo abaixo, comprimimos os pesos de um classificador MNIST para um tamanho muito menor do que a representação de ponto flutuante dele, ainda retendo a exatidão da classificação. Isso é feito em um processo de duas etapas, baseado no artigo Scalable Model Compression by Entropy Penalized Reparameterization (Compressão de modelo escalável pela reparametrização de entropia penalizada):

• Treinando um modelo "compressível" com uma penalidade entrópica explícita durante o treinamento, o que incentiva a compressibilidade dos parâmetros do modelo. O peso dessa penalidade, $\lambda$, permite o controle contínuo do trade-off entre o tamanho do modelo comprimido e a exatidão dele.

• Codificando o modelo compressível em um modelo comprimido usando um esquema de codificação compatível com a penalidade, o que significa que a penalidade é um bom indicador do tamanho do modelo. Isso garante que o método não exija várias iterações de treinamento, compressão e retreinamento do modelo para ajuste.

Esse método é estritamente relacionado ao tamanho do modelo comprimido, e não à complexidade computacional. Ele pode ser combinado com uma técnica como o pruning de modelos para reduzir o tamanho e a complexibilidade.

Exemplo de resultados da compressão em vários modelos:

Modelo (dataset)	Tamanho do modelo	Razão de comp.	Erro top-1 de comp. (não comp.)
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%)

As aplicações incluem:

• Implantar/transmitir modelos para dispositivos de borda em grande escala, economizando largura de banda no trânsito.

• Comunicar o estado de modelo global aos clientes no aprendizado federado. A arquitetura do modelo (número de unidades ocultas etc.) não muda em relação ao modelo inicial, e os clientes podem continuar aprendendo com o modelo descomprimido.

• Realizar a inferência em clientes de memória extremamente limitada. Durante a inferência, os pesos de cada camada podem ser descomprimidos em sequência e descartados logo após as ativações serem computadas.

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

Defina e treine um classificador MNIST básico

Para comprimir camadas convolucionais e densas de maneira eficaz, precisamos definir classes de camadas personalizadas. Elas são análogas às camadas em tf.keras.layers, mas vamos dividi-las depois em subclasses para implementar com eficácia a Reparametrização de Entropia Penalizada (EPR). Para esse fim, também adicionamos um construtor de cópia.

Primeiro, definimos uma camada densa padrão:

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)

E, de maneira similar, uma camada 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 com a compressão do modelo, vamos conferir se podemos treinar um classificador regular.

Defina a arquitetura do 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")

Carregue os dados de treinamento:

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)

Por fim, treine o 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}")

Sucesso! O modelo foi treinado e alcançou uma exatidão de mais de 98% no dataset de validação em 5 épocas.

Treine um classificador compressível

A Reparametrização de Entropia Penalizada (EPR) tem dois ingredientes principais:

• Aplicar uma penalidade aos pesos do modelo durante o treinamento que corresponde à entropia deles em um modelo probabilístico, que é compatível com o esquema de codificação dos pesos. Abaixo, definimos um Regularizer do Keras que implementa essa penalidade.

• Reparametrizar os pesos, ou seja, trazê-los para uma representação latente mais compressível (gera um melhor trade-off entre a compressibilidade e o desempenho do modelo). Para kernels convolucionais, mostrou-se que o domínio Fourier é uma boa representação. Para outros parâmetros, o exemplo abaixo simplesmente usa a quantização escalar (arredondamento) com um tamanho de passo de quantização variável.

Primeiro, defina a penalidade.

O exemplo abaixo usa um modelo probabilístico/código implementado na classe tfc.PowerLawEntropyModel, inspirado no artigo Optimizing the Communication-Accuracy Trade-off in Federated Learning with Rate-Distortion Theory (Otimizando o trade-off entre comunicação e exatidão no aprendizado federado com a teoria taxa-distorção). A penalidade é definida como: $$\log \Bigl(\frac {|x| + \alpha} \alpha\Bigr),$$ em que $x$ é um elemento do parâmetro do modelo ou a representação latente dele, e $\alpha$ é uma pequena constante para a estabilidade numérica em torno de valores 0.

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

A penalidade é efetivamente uma perda de regularização (às vezes chamada de "perda de peso"). O fato de ser côncava com uma cúspide no zero incentiva a dispersão de peso. O esquema de codificação aplicado para a compressão dos pesos, um código Elias gamma, produz códigos de comprimento de $1 + \lfloor \log_2 |x| \rfloor$ bits para a magnitude do elemento. Ou seja, é compatível com a penalidade, e a aplicação da penalidade minimiza o comprimento de código esperado.

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

Segundo, defina subclasses de CustomDense e CustomConv2D, que têm a seguinte funcionalidade adicional:

• Elas pegam uma instância do regularizer acima e a aplicam aos kernels e biases durante o treinamento.

• Elas definem o kernel e bias como uma @property, que realiza a quantização com gradientes diretos quando as variáveis são acessadas. Isso reflete exatamente a computação realizada mais tarde no modelo comprimido.

• Elas definem variáveis log_step adicionais, que representam o logaritmo do tamanho do passo de quantização. Quanto mais espessa a quantização, menor será o tamanho do modelo e também a exatidão. Os tamanhos do passo de quantização são treináveis para cada parâmetro do modelo, então a otimização na função de perda penalizada determinará o melhor tamanho.

O passo de quantização é definido da seguinte maneira:

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

Com isso, podemos definir a camada 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")

A camada convolucional é análoga. Além disso, o kernel de convolução é armazenado como a transformada discreta de Fourier de valor real (RDFT) quando o kernel é definido, e a transformada é invertida sempre que o kernel é usado. Como os componentes de frequência diferente do kernel tendem a ser mais ou menos compressíveis, cada um recebe seu próprio tamanho de passo de quantização.

Defina a transformada de Fourier e o inverso dela da seguinte maneira:

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

Com isso, defina a camada 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 um modelo classificador com a mesma arquitetura que acima, mas usando estas camadas 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)

E treine o modelo:

penalized_accuracy = train_model(
    compressible_classifier, training_dataset, validation_dataset)

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

O modelo compressível alcançou uma exatidão semelhante ao classificador básico.

No entanto, o modelo ainda não foi comprimido. Para fazer isso, precisamos definir outro conjunto de subclasses que armazenam os kernels e biases na forma comprimida — como uma sequência de bits.

Comprima o classificador

As subclasses de CustomDense e CustomConv2D definidas abaixo convertem os pesos de uma camada densa compressível em strings binárias. Além disso, elas armazenam o logaritmo do tamanho de passo da quantização com metade da exatidão para poupar espaço. Quando o kernel ou bias é acessado pela @property, eles são descomprimidos da representação da string e desquantizados.

Primeiro, defina as funções para comprimir e descomprimir um parâmetro de 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

Com elas, 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)

A classe de camada convolucional é análoga à acima.

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 o modelo compressível em comprimido, podemos usar de maneira conveniente a função clone_model. compress_layer converte qualquer camada compressível em uma comprimida e simplesmente passa por qualquer outro tipo de camada (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)

Agora, vamos validar se o modelo comprimido ainda tem o desempenho esperado:

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

A exatidão da classificação do modelo comprimido é idêntica à alcançada durante o treinamento!

Além disso, o tamanho dos pesos do modelo comprimido é bem menor do que o tamanho do 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")

O armazenamento de modelos no disco exige um pouco de sobrecarga para armazenar a arquitetura do modelo, os grafos de função, etc.

Métodos de compressão sem perda como o ZIP são bons para comprimir esse tipo de dados, mas não os próprios pesos. Por isso, a EPR ainda tem um benefício significativo ao contar o tamanho do modelo incluindo essa sobrecarga, após aplicar a compressão de 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")

Efeito de regularização e trade-off entre tamanho e exatidão

Acima, o hiperparâmetro $\lambda$ foi definido como 2 (normalizado pelo número de parâmetros no modelo). Conforme aumentamos o $\lambda$, os pesos do modelo são penalizados cada vez mais pela compressibilidade.

Para valores baixos, a penalidade pode agir como um regularizer de peso. Na verdade, ela tem um efeito benéfico no desempenho de generalização do classificador e pode levar a uma exatidão ligeiramente mais alta no dataset de validação:

#@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 mais altos, vemos um tamanho de modelo cada vez menor, mas também uma exatidão que diminui gradualmente. Para ver isso, vamos treinar alguns modelos e plotar o tamanho em comparação com a exatidão:

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

O ideal é que o plot mostre um trade-off entre o tamanho e a exatidão em formato de cotovelo, mas é normal que as métricas de exatidão tenham algum ruído. Dependendo da inicialização, a curva pode exibir algumas falhas.

Devido ao efeito de regularização, o modelo comprimido de EPR tem maior exatidão no dataset de teste que o modelo original para valores pequenos de $\lambda$. O modelo comprimido de EPR também é diversas vezes menor, mesmo se compararmos os tamanhos após compressão de ZIP adicional.

Descomprima o classificador

CompressedDense e CompressedConv2D descomprimem seus pesos a cada passo para frente. Isso faz com que sejam ideais para dispositivos com memória limitada, mas a descompressão pode ser computacionalmente cara, especialmente para tamanhos de lotes pequenos.

Para descomprimir o modelo uma vez e usá-lo para treinamento ou inferência adicional, podemos convertê-lo de volta para um modelo usando camadas regulares ou compressíveis. Isso pode ser útil na implantação de modelos ou em casos de aprendizado federado.

Primeiro, ao converter de volta para um modelo básico, podemos realizar a inferência e/ou continuar com o treinamento regular sem uma penalidade de compressão:

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

Observe que a exatidão da validação cai após o treinamento para uma época adicional, já que o treinamento é feito sem regularização.

Como alternativa, podemos converter o modelo de volta para o "compressível", para inferência e/ou treinamento adicional com uma penalidade de compressão:

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

Aqui, a exatidão melhora após o treinamento para uma época adicional.