TensorFlow对神经网络模型剪枝
最近在学习模型压缩的方法,了解到TF有自带的模型优化工具tensorflow_model_optimization,所以用 官方示例来学习该工具是怎么进行模型剪枝的。
1 安装
首先安装TF模型优化包tensorflow-model-optimization
,以下用tfmot
简称
pip install -q tensorflow-model-optimization
2 导入模块
用jupyter运行以下代码
import tempfile
import os
import tensorflow as tf
import numpy as np
from tensorflow import keras
3 训练一个基准模型
对MNIST数据集用tf.keras正常训练一个模型(此时还未剪枝)
# Load MNIST dataset
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()
# Normalize the input image so that each pixel value is between 0 to 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')
tf.keras.models.save_model(model, keras_file, include_optimizer=False)
print(