tensorflow / tensorflow/model-optimization
pruning cannnot reduce .h5 model size
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.6k
- Forks
- 349
- Avg merge
- 3d 2h
- Merged PRs (30d)
- 1
Description
Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is here. An issue for anything not supported should be a feature request.
Describe the bug
「prune_low_magnitude」 cannot reduce the size of tflite or model. I run the pruning with keras example.
https://www.tensorflow.org/model_optimization/guide/pruning/pruning_with_keras
I want to obtain smaller size of .h5 file by using pruning.
Actually, I could reduce model size via gzip. However, the output of gzip compression is 「.zip」 file, so I can't do inference with this .zip file.
How can I get compressed .h5 model?? (Are there any other compression method for create the compressed .h5 model??)
System information
TensorFlow version (installed from source or binary):2.5.0
TensorFlow Model Optimization version (installed from source or binary):0.5.0
Python version: 3.7.10
Describe the expected behavior
The size of pruned .h5 model is smaller the base .h5 model.
Describe the current behavior
model1.h5 : base model
model1.h5 : pruned model
Model size is same between base and pruned model.
Code to reproduce the issue
import tempfile
import os
import numpy as np
import pandas as pd
from datetime import datetime
import random
import tensorflow as tf
from tensorflow import keras
# import urllib.request
from keras.datasets import mnist
import gzip
import pickle as cPickle
import sys
output_dir = ('temp/')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# Load MNIST dataset
f = gzip.open('mnist.pkl.gz', 'rb')
if sys.version_info < (3,):
data = cPickle.load(f)
else:
data = cPickle.load(f, encoding='bytes')
f.close()
(train_images, train_labels), (test_images, test_labels) = data
# Normalize the input image so that each pixel value is between 0 and 1.
train_images = train_images / 255.0
test_images = test_images / 255.0
# Define the model architecture.
model = keras.Sequential([
keras.layers.InputLayer(input_shape=(28, 28)),
keras.layers.Reshape(target_shape=(28, 28, 1)),
keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),
keras.layers.MaxPooling2D(pool_size=(2, 2)),
keras.layers.Flatten(),
keras.layers.Dense(10)
])
# Train the digit classification model
model.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
model.fit(
train_images,
train_labels,
epochs=4,
validation_split=0.1,
)
_, baseline_model_accuracy = model.evaluate(
test_images, test_labels, verbose=0)
print('Baseline test accuracy:', baseline_model_accuracy)
_, keras_file = tempfile.mkstemp('.h5')
keras_file = ('temp/model1.h5') #! base model 1
tf.keras.models.save_model(model, keras_file, include_optimizer=False)
print('Saved baseline model to:', keras_file)
import tensorflow_model_optimization as tfmot
prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude
# Compute end step to finish pruning after 2 epochs.
batch_size = 128
epochs = 2
validation_split = 0.1 # 10% of training set will be used for validation set.
num_images = train_images.shape[0] * (1 - validation_split)
end_step = np.ceil(num_images / batch_size).astype(np.int32) * epochs
# Define model for pruning.
pruning_params = {
'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.50,
final_sparsity=0.80,
begin_step=0,
end_step=end_step)
}
model_for_pruning = prune_low_magnitude(model, **pruning_params)
# `prune_low_magnitude` requires a recompile.
model_for_pruning.compile(optimizer='adam',
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
model_for_pruning.summary()
logdir = tempfile.mkdtemp()
callbacks = [
tfmot.sparsity.keras.UpdatePruningStep(),
tfmot.sparsity.keras.PruningSummaries(log_dir=logdir),
]
model_for_pruning.fit(train_images, train_labels, batch_size=batch_size, epochs=epochs, validation_split=validation_split, callbacks=callbacks)
_, model_for_pruning_accuracy = model_for_pruning.evaluate(
test_images, test_labels, verbose=0)
print('Baseline test accuracy:', baseline_model_accuracy)
print('Pruned test accuracy:', model_for_pruning_accuracy)
model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning) # ・・・(※)
# model_for_export=model_for_pruning ・・・(★)
_, pruned_keras_file = tempfile.mkstemp('.h5')
pruned_keras_file = ('temp/model2.h5') #! pruned model 1
tf.keras.models.save_model(model_for_export, pruned_keras_file, include_optimizer=False)
print('Saved pruned Keras model to:', pruned_keras_file)
converter = tf.lite.TFLiteConverter.from_keras_model(model_for_export)
pruned_tflite_model = converter.convert()
_, pruned_tflite_file = tempfile.mkstemp('.tflite')
pruned_tflite_file = ('temp/model3.tflite')
with open(pruned_tflite_file, 'wb') as f:
f.write(pruned_tflite_model)
print('Saved pruned TFLite model to:', pruned_tflite_file)
def get_gzipped_model_size(file):
# Returns size of gzipped model, in bytes.
import os
import zipfile
_, zipped_file = tempfile.mkstemp('.zip')
with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:
f.write(file)
return os.path.getsize(zipped_file)
print("Size of gzipped baseline Keras model: %.2f bytes" % (get_gzipped_model_size(keras_file)))
print("Size of gzipped pruned Keras model: %.2f bytes" % (get_gzipped_model_size(pruned_keras_file)))
print("Size of gzipped pruned TFlite model: %.2f bytes" % (get_gzipped_model_size(pruned_tflite_file)))
converter = tf.lite.TFLiteConverter.from_keras_model(model_for_export)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
quantized_and_pruned_tflite_model = converter.convert()
_, quantized_and_pruned_tflite_file = tempfile.mkstemp('.tflite')
quantized_and_pruned_tflite_file = ('temp/model4.tflite')
with open(quantized_and_pruned_tflite_file, 'wb') as f:
f.write(quantized_and_pruned_tflite_model)
print('Saved quantized and pruned TFLite model to:', quantized_and_pruned_tflite_file)
print("Size of gzipped baseline Keras model: %.2f bytes" % (get_gzipped_model_size(keras_file)))
print("Size of gzipped pruned and quantized TFlite model: %.2f bytes" % (get_gzipped_model_size(quantized_and_pruned_tflite_file)))
Additional context
If I use (★) code instead of (※) code, the pruned model size increased.
(I thought pruned model size reduced ,because strip_pruning restore the original model.
https://www.tensorflow.org/model_optimization/api_docs/python/tfmot/sparsity/keras/strip_pruning?hl=ja)
I would be more appreciate if you also tell the reason about this.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the pruning_with_keras example and the calls to prune_low_magnitude, strip_pruning, and TFLiteConverter in the report. Reproduce the size comparisons for the baseline, stripped, and quantized models using the listed TensorFlow and Model Optimization versions. Done means establishing whether the reported size expectation is supported and documenting or correcting the behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100