TensorFlow 2量化感知训练(QAT)与tf.GradientTape

Question

有人能指出参考，在那里你可以学习如何执行量化感知训练(QAT)与tf.GradientTape上的TensorFlow 2？

我只看到这是用tf.keras API完成的。我不使用tf. keras，我总是用tf.GradientTape构建定制培训，为培训过程提供更多的控制。现在我需要量化一个模型，但是我只看到关于如何使用tf. keras API来实现它的参考资料。

Answer 1

在官方的例子这里中，他们展示了model. fit的QAT训练。下面是使用进行量化感知训练的演示。但是为了完整的参考，让我们在这里两者都做。

基本模型训练。这是直接来自官方医生的。欲知更多详情，请查收。

import os
import tensorflow as tf
from tensorflow import keras
import tensorflow_model_optimization as tfmot

# 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.summary()
model.fit(
  train_images,
  train_labels,
  epochs=1,
  validation_split=0.1,
)

10ms/step - loss: 0.5411 - accuracy: 0.8507 - val_loss: 0.1142 - val_accuracy: 0.9705
<tensorflow.python.keras.callbacks.History at 0x7f9ee970ab90>

QAT .fit.

现在，在基本模型上执行QAT。

# -----------------------
# ------------- Quantization Aware Training -------------
import tensorflow_model_optimization as tfmot

quantize_model = tfmot.quantization.keras.quantize_model
# q_aware stands for for quantization aware.
q_aware_model = quantize_model(model)

# `quantize_model` requires a recompile.
q_aware_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

q_aware_model.summary()
train_images_subset = train_images[0:1000] 
train_labels_subset = train_labels[0:1000]
q_aware_model.fit(train_images_subset, train_labels_subset,
                  batch_size=500, epochs=1, validation_split=0.1)


356ms/step - loss: 0.1431 - accuracy: 0.9629 - val_loss: 0.1626 - val_accuracy: 0.9500
<tensorflow.python.keras.callbacks.History at 0x7f9edf0aef90>

检查性能

_, baseline_model_accuracy = model.evaluate(
    test_images, test_labels, verbose=0)

_, q_aware_model_accuracy = q_aware_model.evaluate(
   test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Quant test accuracy:', q_aware_model_accuracy)

Baseline test accuracy: 0.9660999774932861
Quant test accuracy: 0.9660000205039978

QAT tf.GradientTape().

下面是基础模型上的QAT培训部分。注意，我们还可以在基本模型上执行自定义培训。

batch_size = 500

train_dataset = tf.data.Dataset.from_tensor_slices((train_images_subset,
                                                     train_labels_subset))
train_dataset = train_dataset.batch(batch_size=batch_size, 
                                    drop_remainder=False)

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
optimizer = tf.keras.optimizers.Adam()

for epoch in range(1):
    for x, y in train_dataset:
        with tf.GradientTape() as tape:
            preds = q_aware_model(x, training=True)
            loss = loss_fn(y, preds)
        grads = tape.gradient(loss, q_aware_model.trainable_variables)
        optimizer.apply_gradients(zip(grads, q_aware_model.trainable_variables))
        
_, baseline_model_accuracy = model.evaluate(
    test_images, test_labels, verbose=0)

_, q_aware_model_accuracy = q_aware_model.evaluate(
   test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Quant test accuracy:', q_aware_model_accuracy)

Baseline test accuracy: 0.9660999774932861
Quant test accuracy: 0.9645000100135803