TensorFlow（tfmot）对神经网络模型剪枝

最近在学习模型压缩的方法，了解到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(