Tensorflow 剪枝模型与原始基线模型大小相同

如何解决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

解决方法

在我看来，您似乎错过了实际进行修剪的步骤。如果我们看一下 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)))

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

Why is my pruned model larger than my base model when using Tensorflow's Model Optimization library to prune weights

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

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