Can Quantization Aware Training be performed without using TFLite?

29 views Asked by At

I'm doing Quantization Aware Training. I watch the official example Quantization aware training in Keras example. However, I just want to obtain weights of layer(especially Conv1D) is 8-bits(int8 or float8) instead of float32, and then save into .h5 file, I don't want to use TFLite backend. Is it possible?

Here's my code

import tensorflow as tf
import numpy as np
def build_model():
  model = tf.keras.models.Sequential([
  tf.keras.layers.Conv1D(4, 16, strides= 1,padding='same', activation= 'relu'),
  tf.keras.layers.MaxPooling1D(pool_size=3, strides=2, padding='same'),
  tf.keras.layers.Conv1D(16, 16, strides=1, padding='same', activation='relu'),
  tf.keras.layers.MaxPooling1D(pool_size=3, strides=2, padding='same'),
  tf.keras.layers.Conv1D(32, 16, strides=1, padding='same', activation='relu'),
  tf.keras.layers.AveragePooling1D(pool_size=3, strides=2, padding='same'),
  tf.keras.layers.Conv1D(64, 16, strides=1, padding= 'same', activation='relu'),
  tf.keras.layers.MaxPooling1D(pool_size=3, strides=2, padding='same'),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(128, activation='relu'),
  # tf.keras.layers.Dropout(0.5),
  tf.keras.layers.Dense(15, activation='softmax')
  ])

  model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['acc'])
  return model


model = build_model()

#QAT
import tensorflow_model_optimization as tfmot
from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_configs
NoOpQuantizeConfig = default_8bit_quantize_configs.NoOpQuantizeConfig

class NoOpQuantizeConfig(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 {}

def apply_quantization(layer):
    if isinstance(layer, (tf.keras.layers.Conv1D, tf.keras.layers.MaxPooling1D)):
        return tfmot.quantization.keras.quantize_annotate_layer(layer, quantize_config=NoOpQuantizeConfig())
    else:
        return tfmot.quantization.keras.quantize_annotate_layer(layer)

annotated_adv_model = tf.keras.models.clone_model(
    model, clone_function=apply_quantization,
)

with tf.keras.utils.custom_object_scope({"NoOpQuantizeConfig": NoOpQuantizeConfig}):
    q_aware_model = tfmot.quantization.keras.quantize_model(annotated_adv_model)

q_aware_model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['acc'])
q_aware_model.fit(signal_data, symbol_data, epochs=1, batch_size=128, validation_split=0.3, callbacks=callbacks_list)
q_aware_model.save('/content/drive/MyDrive/qaware_model.h5')
0

There are 0 answers