CODEWITHSUNDEEP
pythontensorflowkeraspruning
LucasStromberg
LucasStromberg

Reputation: 53

Tensorflow pruned model is the same size as original baseline model

I have a baseline TF functional model that I want to prune. I have tried following the code in the documentation, but the size the compressed pruned model is the same size as the compressed baseline model.

(https://www.tensorflow.org/model_optimization/guide/pruning/comprehensive_guide#export_model_with_size_compression)

I don't believe there is anything wrong with my code, so why does this occur?

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

Output:

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

EDIT:

I also tried this with the same model as in the documentation, and the pruned model is still the same size as the baseline:

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

Output:

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

Upvotes: 0

Views: 445

Answers (1)

djvaroli
djvaroli

Reputation: 1383

It looks to me like you are missing the step that actually does the pruning. If we look at the test() function, you set the model up for pruning but never actually prune it. Look at the edits below.

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

You can take a look at the code I have in a question I just asked. I am having a slightly different problem but the code posted there works (at least for some of the cases).

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

You can also check out the tensorflow.sparsity.keras APIs to see some of the other options if you are interested

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

Upvotes: 1

Related Questions