可扩展的模型压缩

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 Github 上查看源代码

下载笔记本

概述

本笔记本展示了如何使用 TensorFlow Compression 压缩模型。

在下面的示例中，我们将 MNIST 分类器的权重压缩到比其浮点表示小得多的大小，同时保持分类准确率。这是通过基于论文 Scalable Model Compression by Entropy Penalized Reparameterization 的两步过程完成的：

• 在训练期间使用显式熵惩罚来训练“可压缩”模型，这鼓励了模型参数的可压缩性。此惩罚的权重 \(\lambda\)，能够持续控制压缩模型大小和其准确率之间的权衡。

• 使用与惩罚相匹配的编码方案将可压缩模型编码为压缩模型，这意味着惩罚是对模型大小的良好预测指标。这确保了该方法不需要多次迭代训练、压缩和重新训练模型以进行微调。

这种方法会严格考虑压缩模型的大小，而不是计算复杂度。它可以与模型剪枝等技术相结合，以减少大小和复杂度。

各种模型的压缩结果示例：

应用包括：

• 大规模部署/广播模型到边缘设备，节省传输带宽。
• 在联合学习中向客户端传达全局模型状态。模型架构（隐藏单元的数量等）相较于初始模型没有变化，客户端可以在解压缩的模型上继续学习。
• 在内存极其有限的客户端上执行推断。在推断过程中，可以按顺序解压缩每一层的权重，并在计算激活后立即丢弃。

设置

通过 pip 安装 TensorFlow Compression。

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

导入库依赖项。

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

2023-11-07 22:10:27.535568: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-11-07 22:10:27.535626: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-11-07 22:10:27.535666: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

定义和训练一个基本的 MNIST 分类器

为了高效压缩密集层和卷积层，我们需要定义自定义层类。这些类似于 tf.keras.layers 下的层，但我们稍后将对它们进行子类化以高效实现熵惩罚重参数化 (EPR)。为此，我们还添加了一个复制构造函数。

首先，我们定义一个标准的密集层：

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)

类似地，定义一个 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)

在继续模型压缩之前，我们来检查一下是否可以成功地训练一个常规分类器。

定义模型架构：

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

2023-11-07 22:10:31.096992: W tensorflow/core/common_runtime/gpu/gpu_device.cc:2211] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

加载训练数据：

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)

最后，训练模型：

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

Epoch 1/5
469/469 [==============================] - 53s 111ms/step - loss: 0.2067 - sparse_categorical_accuracy: 0.9383 - val_loss: 0.0837 - val_sparse_categorical_accuracy: 0.9741
Epoch 2/5
469/469 [==============================] - 52s 110ms/step - loss: 0.0635 - sparse_categorical_accuracy: 0.9807 - val_loss: 0.0649 - val_sparse_categorical_accuracy: 0.9782
Epoch 3/5
469/469 [==============================] - 52s 111ms/step - loss: 0.0436 - sparse_categorical_accuracy: 0.9864 - val_loss: 0.0563 - val_sparse_categorical_accuracy: 0.9820
Epoch 4/5
469/469 [==============================] - 52s 111ms/step - loss: 0.0324 - sparse_categorical_accuracy: 0.9901 - val_loss: 0.0568 - val_sparse_categorical_accuracy: 0.9828
Epoch 5/5
469/469 [==============================] - 52s 111ms/step - loss: 0.0266 - sparse_categorical_accuracy: 0.9916 - val_loss: 0.0600 - val_sparse_categorical_accuracy: 0.9836
Accuracy: 0.9836

成功！该模型训练良好，在 5 个周期内的验证集上的准确率达到了 98% 以上。

训练可压缩分类器

熵惩罚重参数化（EPR）有两个主要组成部分：

• 在训练期间对模型权重施加惩罚，该惩罚对应于概率模型下的熵，并与权重的编码方案相匹配。下面，我们定义一个实现此惩罚的 Keras Regularizer。

• 重新参数化权重，即将它们带入更具可压缩性的潜在表示中（在可压缩性和模型性能之间达成更好的权衡）。对于卷积核，已经证明傅里叶域是一个很好的表示。对于其他参数，以下示例仅使用具有不同量化步长的标量量化（舍入）。

首先，定义惩罚。

下面的示例使用在 tfc.PowerLawEntropyModel 类中实现的代码/概率模型，灵感来自论文 Optimizing the Communication-Accuracy Trade-off in Federated Learning with Rate-Distortion Theory。惩罚定义为：\( \log \Bigl(\frac {|x| + \alpha} \alpha\Bigr)，\) 其中 \(x\) 是模型参数或其潜在表示的一个元素，\(\alpha\) 是一个数值稳定性在 0 附近小常量。

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

这种惩罚实际上是一种正则化损失（有时称为“权重损失”）。它是凹形的，顶点为零，这一事实鼓励权重稀疏。用于压缩权重的编码方案是 Elias gamma 码，它为元素大小产生长度为 \( 1 + \lfloor \log_2 |x| \rfloor \) 比特的编码。也就是说，它与惩罚相匹配，并应用惩罚从而最小化预期的代码长度。

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

其次，定义 CustomDense 和 CustomConv2D 的子类，它们具有以下附加功能：

• 它们接受上述 Regularizer 的一个实例，并将其应用于训练期间的内核和偏差。
• 它们将内核和偏差定义为 @property，每当访问变量时，它们都会使用直通梯度执行量化。这准确地反映了稍后在压缩模型中执行的计算。
• 它们定义了额外的 log_step 变量，代表量化步长的对数。量化越粗，模型越小，但准确率越低。每个模型参数的量化步长都可训练，因此对惩罚损失函数执行优化将确定最佳量化步长。

量化步长定义如下：

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

有了它，我们可以定义密集层：

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

卷积层类似。此外，只要设置了卷积核，就会将卷积核作为其实值离散傅里叶变换 (RDFT) 存储，并且每当使用该核时，变换都会被反转。由于内核的不同频率分量往往或多或少是可压缩的，因此其中的每个分量都被分配了自己的量化步长。

按如下方式定义傅里叶变换及其逆变换：

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

这样，将卷积层定义为：

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

使用与上面相同的架构定义分类器模型，但使用以下修改后的层：

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)

并训练模型：

penalized_accuracy = train_model(
    compressible_classifier, training_dataset, validation_dataset)

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

Epoch 1/5
469/469 [==============================] - 58s 118ms/step - loss: 3.7884 - sparse_categorical_accuracy: 0.9278 - val_loss: 2.1724 - val_sparse_categorical_accuracy: 0.9730
Epoch 2/5
469/469 [==============================] - 55s 118ms/step - loss: 1.6649 - sparse_categorical_accuracy: 0.9767 - val_loss: 1.3018 - val_sparse_categorical_accuracy: 0.9811
Epoch 3/5
469/469 [==============================] - 55s 117ms/step - loss: 1.0736 - sparse_categorical_accuracy: 0.9829 - val_loss: 0.9402 - val_sparse_categorical_accuracy: 0.9830
Epoch 4/5
469/469 [==============================] - 56s 119ms/step - loss: 0.7910 - sparse_categorical_accuracy: 0.9859 - val_loss: 0.8004 - val_sparse_categorical_accuracy: 0.9831
Epoch 5/5
469/469 [==============================] - 56s 119ms/step - loss: 0.6509 - sparse_categorical_accuracy: 0.9872 - val_loss: 0.6408 - val_sparse_categorical_accuracy: 0.9837
Accuracy: 0.9837

可压缩模型已达到与普通分类器相似的准确率。

但是，该模型实际上还没有被压缩。为此，我们定义了另一组子类，它们以压缩形式存储内核和偏差（作为位序列）。

压缩分类器

下面定义的 CustomDense 和 CustomConv2D 的子类将可压缩密集层的权重转换为二进制字符串。此外，它们以半精度存储量化步长的对数以节省空间。每当通过 @property 访问内核或偏差时，它们就会从其字符串表示中解压缩并去量化。

首先，定义函数来压缩和解压缩模型参数：

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

有了这些，我们可以定义 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)

卷积层类与上面类似。

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)

要将可压缩模型转换为压缩模型，我们可以方便地使用 clone_model 函数。compress_layer 可以将任何可压缩层转换为压缩层，并简单地传递给任何其他类型的层（例如 Flatten 等）。

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)

现在，我们来验证压缩模型是否仍按预期执行：

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

79/79 [==============================] - 1s 11ms/step - loss: 0.0000e+00 - sparse_categorical_accuracy: 0.9837
Accuracy of the compressible classifier: 0.9837
Accuracy of the compressed classifier: 0.9837

压缩模型的分类准确率与训练期间达到的分类准确率相同！

此外，压缩后的模型权重的大小远小于原始模型的大小：

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

Size of original model weights: 5024320 bytes
Size of compressed model weights: 18653 bytes
Compression ratio: 269x

将模型存储在磁盘上需要一些开销来存储模型架构、函数图等。

ZIP 等无损压缩方法擅长压缩此类数据，但不擅长压缩权重本身。这就是为什么在应用了 ZIP 压缩之后，当计算模型大小（包括开销）时，EPR 仍然具有显著优势：

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

INFO:tensorflow:Assets written to: /tmp/classifier/assets
INFO:tensorflow:Assets written to: /tmp/classifier/assets
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
Original on-disk size (ZIP compressed): 13900646 bytes
Compressed on-disk size (ZIP compressed): 60810 bytes
Compression ratio: 229x

正则化效果和大小-准确度权衡

上面，\(\lambda\) 超参数被设置为 2（通过模型中的参数数量进行标准化）。随着我们增加 \(\lambda\)，模型权重的可压缩性受到越来越严重的惩罚。

对于较低的值，惩罚可以起到权重调节器的作用。它实际上对分类器的泛化性能有有益的影响，并且可以在验证数据集上产生略高的准确率：

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

Accuracy of the vanilla classifier: 0.9836
Accuracy of the penalized classifier: 0.9837

对于更高的值，我们看到模型大小越来越小，但准确率也在逐渐降低。为了看到这一点，我们来训练几个模型，并绘制它们的大小与准确率之间的关系图：

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)

lambda=2: training...
compressing...
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
evaluating...
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
lambda=5: training...
compressing...
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
evaluating...
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
lambda=10: training...
compressing...
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
evaluating...
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
lambda=20: training...
compressing...
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
evaluating...
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
lambda=50: training...
compressing...
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
INFO:tensorflow:Assets written to: /tmp/compressed_classifier/assets
evaluating...
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.
WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.

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

理想情况下，该图应显示肘形大小-准确率权衡，但准确率指标有些噪声也正常。根据初始化的不同，曲线可能会出现一些曲折。

由于正则化效应，对于较小的 \(\lambda\) 值，EPR 压缩模型在测试集上比原始模型更准确。即使我们比较附加 ZIP 压缩后的大小，EPR 压缩模型也要小很多倍。

解压缩分类器

CompressedDense 和 CompressedConv2D 在每次前向传递时会解压缩它们的权重。这使得它们非常适合内存有限的设备，但解压缩的计算成本可能很高，尤其是对于小批次。

要将模型解压缩一次，并将其用于进一步的训练或推断，我们可以使用常规层或可压缩层将其转换回模型。这在模型部署或联合学习场景中很有用。

首先，转换回普通模型，我们可以进行推断，和/或继续进行常规训练，而不会有压缩惩罚：

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

469/469 [==============================] - 53s 111ms/step - loss: 0.0828 - sparse_categorical_accuracy: 0.9750 - val_loss: 0.0665 - val_sparse_categorical_accuracy: 0.9792
Accuracy of the compressed classifier: 0.9837
Accuracy of the decompressed classifier after one more epoch of training: 0.9792

请注意，在训练额外的周期后验证准确率会下降，因为训练是在没有正则化的情况下完成的。

或者，我们可以将模型转换回“可压缩”模型，以进行推断和/或进一步训练，并带有压缩惩罚：

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

469/469 [==============================] - 58s 118ms/step - loss: 0.7861 - sparse_categorical_accuracy: 0.9897 - val_loss: 0.7304 - val_sparse_categorical_accuracy: 0.9869
Accuracy of the compressed classifier: 0.9837
Accuracy of the decompressed classifier after one more epoch of training: 0.9869

在这里，在训练一个额外的周期后，准确率会提高。