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

我想知道目前在Tensorflow 2中模拟量化感知训练期间BatchNorm折叠的可用选项是什么。 Tensorflow 1有 tf.contrib.quantize.create_training_graph 函数，该函数将FakeQuantization层插入图中，并负责模拟批量归一化折叠（根据这个 白皮书).

Tensorflow 2有一个 教程 如何在其最近通过的《公约》中使用量化的方法。tf.keras API，但他们没有提到任何关于批处理正常化的内容。我尝试了下面这个带有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，用老方法定义一个图？

# 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)

StackOverflow

收藏 分享 链接