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.

확장 가능한 모델 압축

TensorFlow.org에서 보기

Google Colab에서 실행하기

GitHub에서 소스 보기

노트북 다운로드하기

개요

이 노트북은 TensorFlow Compression을 사용하여 모델을 압축하는 방법을 보여줍니다.

아래 예에서는 분류 정확도를 유지하면서 MNIST 분류자의 가중치를 부동 소수점 표현보다 훨씬 작은 크기로 압축합니다. 이 작업은 Scalable Model Compression by Entropy Penalized Reparameterization 논문에 기초한 두 단계 프로세스로 수행됩니다.

• 훈련 중에 명시적 엔트로피 페널티로 "압축 가능" 모델을 훈련하여 모델 매개변수의 압축성을 유도합니다. 이 페널티에 대한 가중치 $\lambda$를 통해 압축된 모델 크기와 정확도 간의 균형을 지속적으로 제어할 수 있습니다.

• 페널티와 일치하는 코딩 방식을 사용하여 압축 가능한 모델을 압축된 모델로 인코딩합니다. 즉, 페널티가 모델 크기에 대한 좋은 예측 변수입니다. 이렇게 하면 메서드가 미세 조정을 위해 모델 훈련, 압축 및 재훈련을 여러 번 반복할 필요가 없습니다.

이 방법은 계산 복잡성이 아니라 압축된 모델 크기에만 엄격히 관련됩니다. 크기와 복잡성을 줄이기 위해 모델 정리(model pruning)와 같은 기술과 결합할 수 있습니다.

다양한 모델에서 압축 결과의 예:

모델(데이터세트)	모델 크기	압축 비율	상위 1 오차 보정(비압축)
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%)

애플리케이션은 다음을 포함합니다.

• 모델을 대규모로 에지 장치에 배포/브로드캐스팅하여 전송 대역폭을 절약합니다.

• 페더레이션 학습에서 클라이언트에게 글로벌 모델 상태를 전달합니다. 모델 아키텍처(숨겨진 유닛의 수 등)는 초기 모델에서 변경되지 않았으며 클라이언트는 압축 해제된 모델에 대해 학습을 계속할 수 있습니다.

• 메모리가 크게 제한된 클라이언트에서 추론을 수행합니다. 추론하는 동안 각 레이어의 가중치는 순차적으로 압축 해제되고 활성화가 계산된 직후 버릴 수 있습니다.

설정

pip를 통해 Tensorflow Compression을 설치합니다.

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

라이브러리 종속성을 가져옵니다.

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

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

훈련 데이터를 로드합니다.

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

성공입니다! 모델이 잘 훈련되었고 5번의 epoch 내에서 검증 세트에 대해 98% 이상의 정확도에 도달했습니다.

압축 가능한 분류자 훈련하기

EPR(Entropy Penalized Reparameterization)에는 두 가지 주요 구성 요소가 있습니다.

• 가중치의 인코딩 방식과 일치하는 확률 모델에서 엔트로피에 해당하는 모델 가중치에 페널티를 훈련 중에 적용합니다. 아래에서 이 페널티를 구현하는 Keras Regularizer를 정의합니다.

• 가중치를 다시 매개변수화합니다. 즉, 더 압축 가능한 잠재 표현으로 가져옵니다(압축성과 모델 성능 사이에 더 나은 균형을 제공함). 컨볼루션 커널의 경우 퓨리에 도메인이 좋은 표현인 것으로 나타났습니다. 다른 매개변수의 경우 아래 예는 단순히 다양한 양자화 단계 크기로 스칼라 양자화(반올림)를 사용합니다.

먼저 페널티를 정의합니다.

아래 예제는 tfc.PowerLawEntropyModel 클래스에서 구현된 코드/확률 모델을 사용합니다. 이 모델은 Optimizing 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(_));

페널티는 사실상 정규화 손실입니다(때로 "가중치 손실"이라고도 함). 이것이 0에 첨두를 둔 오목한 형태라는 사실은 가중치 희소성을 높입니다. 가중치 압축에 적용된 코딩 체계인 Elias 감마 코드는 요소의 크기에 대해 $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}")

압축 가능한 모델이 일반 분류자와 유사한 정확도에 도달했습니다.

그러나 이 모델은 아직 실제로 압축되지 않았습니다. 이를 위해 우리는 커널과 바이어스를 압축된 형태로 저장하는 또 다른 하위 클래스 세트를 비트 시퀀스로 정의합니다.

분류자 압축하기

아래에 정의된 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}")

압축된 모델의 분류 정확도는 훈련 중에 얻은 것과 동일합니다!

또한 압축된 모델 가중치의 크기는 원래 모델 크기보다 훨씬 작습니다.

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

모델을 디스크에 저장하려면 모델 아키텍처, 함수 그래프 등을 저장하는 데 약간의 오버헤드가 필요합니다.

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

정규화 효과와 크기-정확도 절충

위에서 $\lambda$ 하이퍼파라미터는 2로 설정되었습니다(모델의 매개변수 수로 정규화됨). $\lambda$가 증가함에 따라 모델 가중치는 압축성에 대해 점점 더 많은 패널티를 받습니다.

값이 낮으면 페널티가 가중치 regularizer처럼 작동할 수 있습니다. 이것은 실제로 분류자의 일반화 성능에 유익한 영향을 미치며 유효성 검사 데이터세트에서 약간 더 높은 정확도로 이어질 수 있습니다.

#@title

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

값이 높을수록 모델 크기가 점점 작아지지만 정확도도 점차 감소합니다. 이를 확인하기 위해 몇 가지 모델을 훈련하고 크기 대 정확도 플롯을 나타내 보겠습니다.

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

이 플롯은 이상적으로 팔꿈치 모양의 크기-정확도 절충을 보여야 하지만 정확도 메트릭에 다소 노이즈가 있는 것은 정상입니다. 초기화에 따라 곡선이 약간 꼬일 수 있습니다.

정규화 효과로 인해 EPR 압축 모델은 $\lambda$의 작은 값에 대해 원래 모델보다 테스트 세트에서 더 정확합니다. 추가 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}")

훈련이 정규화 없이 수행되기 때문에 추가 epoch에 대한 훈련 후에 유효성 검사 정확도가 떨어집니다.

또는 압축 페널티로 추론 및/또는 추가 훈련하기 위해 모델을 "압축 가능한" 모델로 다시 변환할 수 있습니다.

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

여기서 추가 epoch에 대한 훈련 후 정확도가 향상됩니다.