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