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

Question

我想知道在 Tensorflow 2 的量化感知训练期间模拟 BatchNorm 折叠的当前可用选项是什么。 Tensorflow 1 具有tf.contrib.quantize.create_training_graph将 FakeQuantization 层插入图中并负责模拟批量归一化折叠的功能（根据本白皮书）。

Tensorflow 2 有一个关于如何在他们最近采用的API 中使用量化的教程tf.keras，但他们没有提到关于批量标准化的任何内容。我使用 BatchNorm 层尝试了以下简单示例：

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)

然而，它给出了以下例外：

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.

这表明 TF 不知道如何处理它。

我还看到了这个相关主题，它们适用tf.contrib.quantize.create_training_graph于 keras 构建的模型。然而，他们不使用 BatchNorm 层，所以我不确定这会起作用。

那么在 TF2 中使用这个 BatchNorm 折叠功能有哪些选择呢？这可以从 keras API 完成，还是应该切换回 TensorFlow 1 API 并以旧方式定义图形？

Answer 1

如果您在激活之前添加 BatchNormalization，则不会出现量化问题。注意：仅当该层正好位于 Conv2D 层之后时，BatchNormalization 才支持量化。 https://www.tensorflow.org/model_optimization/guide/quantization/training

# Change
l.Conv2D(64, 5, padding='same', activation='relu'),
l.BatchNormalization(),    # BN!
# with this
l.Conv2D(64, 5, padding='same'),
l.BatchNormalization(),
l.Activation('relu'),

#Other way of declaring the same
o = (Conv2D(512, (3, 3), padding='valid' , data_format=IMAGE_ORDERING))(o)
o = (BatchNormalization())(o)
o = Activation('relu')(o)