时间戳

首页 > 解决方案 > Tensorflow 修剪模型与原始基线模型大小相同

python - Tensorflow 修剪模型与原始基线模型大小相同

问题描述

我有一个要修剪的基线 TF 功能模型。我尝试按照文档中的代码进行操作，但压缩修剪模型的大小与压缩基线模型的大小相同。

（https://www.tensorflow.org/model_optimization/guide/pruning/comprehensive_guide#export_model_with_size_compression）

我不相信我的代码有什么问题，那么为什么会发生这种情况呢？

def get_gzipped_model_size(model):
  # Returns size of gzipped model, in bytes.
  import os
  import zipfile

  _, keras_file = tempfile.mkstemp('.h5')
  model.save(keras_file, include_optimizer=False)

  _, zipped_file = tempfile.mkstemp('.zip')
  with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:
    f.write(keras_file)

  return os.path.getsize(zipped_file)


def test():
    model = keras.models.load_model('models/cifar10/baselines/convnet_small')
    model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model)

    model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

    print("Size of gzipped baseline model: %.2f bytes" % (get_gzipped_model_size(model)))
    print("Size of gzipped pruned model without stripping: %.2f bytes" % (get_gzipped_model_size(model_for_pruning)))
    print("Size of gzipped pruned model with stripping: %.2f bytes" % (get_gzipped_model_size(model_for_export)))


if __name__ == "__main__":
    test()

输出：

Size of gzipped baseline model: 604286.00 bytes

Size of gzipped pruned model without stripping: 610750.00 bytes

Size of gzipped pruned model with stripping: 604287.00 bytes

编辑：

我还使用与文档中相同的模型进行了尝试，修剪后的模型仍然与基线大小相同：

input_shape = [20]
x_train = np.random.randn(1, 20).astype(np.float32)
y_train = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=20)


def setup_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(20, input_shape=input_shape),
      tf.keras.layers.Flatten()
  ])
  return model

def setup_pretrained_weights():
  model = setup_model()

  model.compile(
      loss=tf.keras.losses.categorical_crossentropy,
      optimizer='adam',
      metrics=['accuracy']
  )

  model.fit(x_train, y_train)

  _, pretrained_weights = tempfile.mkstemp('.tf')

  model.save_weights(pretrained_weights)

  return pretrained_weights


setup_model()
pretrained_weights = setup_pretrained_weights()

输出：

Size of gzipped baseline model: 2910.00 bytes
Size of gzipped pruned model without stripping: 3333.00 bytes
Size of gzipped pruned model with stripping: 2910.00 bytes

标签： pythontensorflowkeraspruning

解决方案

在我看来，您似乎错过了实际进行修剪的步骤。如果我们查看test()函数，您会设置模型以进行修剪，但实际上从未修剪过它。看看下面的编辑。

import tensorflow_model_optimization as tfmot

def test():
    model = keras.models.load_model('models/cifar10/baselines/convnet_small')
    pruning_schedule = tfmot.sparsity.keras.ConstantSparsity(
                         target_sparsity=0.95, 
                         begin_step=0, 
                         end_step=-1, 
                         frequency=100
                        )

    callbacks = [tfmot.sparsity.keras.UpdatePruningStep()]
    model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model, pruning_schedule=pruning_schedule)
    model_for_pruning.compile(optimizer="adam", loss="some-loss")
    model_for_pruning.fit(X, y, epochs=2, callbacks=callbacks)
    model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)


    print("Size of gzipped baseline model: %.2f bytes" %(get_gzipped_model_size(model)))
    print("Size of gzipped pruned model without stripping: %.2f bytes" % (get_gzipped_model_size(model_for_pruning)))
    print("Size of gzipped pruned model with stripping: %.2f bytes" % (get_gzipped_model_size(model_for_export)))

你可以看看我刚才问的问题中的代码。我有一个稍微不同的问题，但那里发布的代码有效（至少在某些情况下）。

为什么使用 Tensorflow 的模型优化库修剪权重时，修剪后的模型比基本模型大

如果您有兴趣，还可以查看tensorflow.sparsity.kerasAPI 以查看其他一些选项

https://www.tensorflow.org/model_optimization/api_docs/python/tfmot/sparsity/keras

推荐阅读