Обучение квантованию в TensorFlow версии 2 и сворачивании BatchNorm

Question

Мне интересно, каковы текущие доступные параметры для имитации сворачивания BatchNorm во время обучения квантованию в Tensorflow 2. Tensorflow 1 имеет tf.contrib.quantize.create_training_graphфункция, которая вставляет слои FakeQuantization в график и заботится о имитации пакетной нормализации сворачивания (согласно этому техническому документу).

В Tensorflow 2 есть руководство о том, как использовать квантование в недавно принятыхtf.kerasAPI, но они ничего не упоминают о пакетной нормализации. Я попробовал следующий простой пример со слоем 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на построенной керас модели. Однако они не используют слои BatchNorm, поэтому я не уверен, что это сработает.

Итак, каковы варианты использования этой функции сворачивания BatchNorm в TF2? Можно ли это сделать с помощью APIkeras или мне следует вернуться к API TensorFlow 1 и определить график по-старому?

Answer 1

Если вы добавите BatchNormalization перед активацией, у вас не будет проблем с квантованием. Примечание. Квантование поддерживается в BatchNormalization только в том случае, если этот слой находится точно после уровня Conv2D. 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)

Answer 2

Следует применить аннотацию квантования, как в инструкции. Думаю, теперь вы можете вызвать BatchNorm вот так:

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 []

def get_config(self):
    return {}

Если вы по-прежнему хотите квантовать слой, измените возвращение get_weights_and_quantizers на return [(layer.weights[i], LastValueQuantizer(num_bits=8, symmetric=True, narrow_range=False, per_axis=False)) for i in range(2)]. Затем верните квантователи на гамму, бета,... в соответствии с индексами списка возврата выше в set_quantize_weights

Результат будет примерно таким (RESNET50):