我想知道在 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 并以旧方式定义图形?