TensorFlow 版本 2 和 BatchNorm 折叠中的量化感知训练

问题描述

我想知道在 Tensorflow 2 中量化感知训练期间模拟 BatchNorm 折叠的当前可用选项是什么.Tensorflow 1 具有 tf.contrib.quantize.create_training_graph 函数，该函数将 FakeQuantization 层插入到图并负责模拟批量归一化折叠(根据此

I'm wondering what the current available options are for simulating BatchNorm folding during quantization aware training in Tensorflow 2. Tensorflow 1 has the tf.contrib.quantize.create_training_graph function which inserts FakeQuantization layers into the graph and takes care of simulating batch normalization folding (according to this white paper).

Tensorflow 2 has a tutorial on how to use quantization in their recently adopted tf.keras API, but they don't mention anything about batch normalization. I tried the following simple example with a BatchNorm layer:

import tensorflow_model_optimization as tfmo

model = tf.keras.Sequential([
      l.Conv2D(32, 5, padding='same', activation='relu', input_shape=input_shape),
      l.MaxPooling2D((2, 2), (2, 2), padding='same'),
      l.Conv2D(64, 5, padding='same', activation='relu'),
      l.BatchNormalization(),    # BN!
      l.MaxPooling2D((2, 2), (2, 2), padding='same'),
      l.Flatten(),
      l.Dense(1024, activation='relu'),
      l.Dropout(0.4),
      l.Dense(num_classes),
      l.Softmax(),
])
model = tfmo.quantization.keras.quantize_model(model)

It however gives the following exception:

RuntimeError: Layer batch_normalization:<class 'tensorflow.python.keras.layers.normalization.BatchNormalization'> is not supported. You can quantize this layer by passing a `tfmot.quantization.keras.QuantizeConfig` instance to the `quantize_annotate_layer` API.

which indicates that TF does not know what to do with it.

I also saw this related topic where they apply tf.contrib.quantize.create_training_graph on a keras constructed model. They however don't use BatchNorm layers, so I'm not sure this will work.

So what are the options for using this BatchNorm folding feature in TF2? Can this be done from the keras API, or should I switch back to the TensorFlow 1 API and define a graph the old way?

解决方案

You should apply the quantization annotation as in the instruction. I think you can call the BatchNorm now like this:

class DefaultBNQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):

def get_weights_and_quantizers(self, layer):
    return []

def get_activations_and_quantizers(self, layer):
    return []

def set_quantize_weights(self, layer, quantize_weights):
    pass
def set_quantize_activations(self, layer, quantize_activations):
    pass
def get_output_quantizers(self, layer):
    return [tfmot.quantization.keras.quantizers.MovingAverageQuantizer(
    num_bits=8, per_axis=False, symmetric=False, narrow_range=False)]

def get_config(self):
    return {}

If you still want to quantize for the layer, change the return of the get_weights_and_quantizers to return [(layer.weights[i], LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False)) for i in range(2)]. Then set back the quantizers to gamma,beta,... according to the indices of the return list above at set_quantize_weights. However, I am not encouraging this way as it surely harm the accuracy as BN should act as an activation quantization

The result you have would be like this (RESNET50):

这篇关于TensorFlow 版本 2 和 BatchNorm 折叠中的量化感知训练的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！