Quantisation aware training worsens autoencoder performance when quantising the output of the encoder.

[Berk-Bilgin]

Env Info: WSL2 (Ubuntu), Python 3.9.19, tensorflow 2.15.1, qkeras 0.9.0

Hello, I am trying to implement the encoder part of an autoencoder model on an FPGA for data compression. The input to the model is 64 samples of 16-bit signed integers (no preprocessing, I tried it, didn't get good results). I am using qkeras layers only for the encoder, and keeping the decoder's layers as keras layers to benefit from the higher precision. When I tried to train the autoencoder I have given below, I saw that the error (custom loss function similar to mse loss) was quite large, so I decided to approach it in steps, quantising the model in parts and training, gradually adding more quantisation with each model. I finally constructed the same quantised model but with the encoder's final layer having a regular keras 'linear' activation instead of quantized_bits(). The result was much better, similar to the performance of the non-quantised model. I decided to pass the weights of this partially quantised model to the model below and its performance was also quite good, only slightly worse. However when I tried to train this model with its weights initialised as I described, the training only increased the loss function. The loss started out good at the start of training and gradually got worse after each epoch. I tried decreasing Adam's learning rate however the increase in the loss function as the training progressed persisted, only slower than before. So basically I am unable to do quantization aware training for the autoencoder when the encoder is fully quantised and would like to know if I am doing something wrong or I should be doing some sort of conversion when going from the fixed point encoder to the floating point decoder. Thanks in advance.

# Define the encoder with quantization
input_ts = keras.layers.Input(shape=(64, 1), name="input_time_series")  # 64, 1

# First Conv1D layer with quantized activations
x = QConv1D(16, 3, 
            kernel_quantizer=quantizers.quantized_bits(16, 0, alpha=1), 
            bias_quantizer=quantizers.quantized_bits(16, 0, alpha=1),
            padding='same')(input_ts)  # 64, 16
x = QActivation(activation= quantizers.quantized_bits(32, 15, alpha=1))(x)  # Quantized activations to 16 bits
# Second Conv1D layer with quantized activations
x = QConv1D(8, 5, 
            kernel_quantizer=quantizers.quantized_bits(16, 0, alpha=1), 
            bias_quantizer=quantizers.quantized_bits(16, 0, alpha=1),
            padding='same')(x)  # 64, 16
x = QActivation(activation= quantizers.quantized_bits(32, 15, alpha=1))(x)
# Third Conv1D layer with quantized activations
x = QConv1D(8, 5, 
            kernel_quantizer=quantizers.quantized_bits(16, 0, alpha=1), 
            bias_quantizer=quantizers.quantized_bits(16, 0, alpha=1),
            padding='same', 
            strides=2)(x)  # 32, 8
x = QActivation(activation= quantizers.quantized_bits(32, 15, alpha=1))(x)
# Fourth Conv1D layer with quantized activations
x = QConv1D(4, 5, 
            kernel_quantizer=quantizers.quantized_bits(16, 0, alpha=1), 
            bias_quantizer=quantizers.quantized_bits(16, 0, alpha=1),
            padding='same', 
            strides=2)(x)  # 16, 4
x = QActivation(activation= quantizers.quantized_bits(32, 15, alpha=1))(x)
# Flatten layer
x = keras.layers.Flatten()(x)  # Flatten for Dense layer

# Dense layer with quantized activations
x = QDense(16, kernel_quantizer=quantizers.quantized_bits(16, 0, alpha=1), 
                 bias_quantizer=quantizers.quantized_bits(16, 0, alpha=1))(x)  # 16
encoded = QActivation(activation= quantizers.quantized_bits(16, 15, alpha=1))(x)

# Define the model
encoder = keras.models.Model(input_ts, encoded, name="encoder")

# Print the model summary
encoder.summary()

# Define the decoder
encoded_input = keras.layers.Input(shape=(16,), name="encoded_input") # 16
x = keras.layers.Dense(16 * 4, activation='linear')(encoded_input) # 16 * 4
x = keras.layers.Reshape((16, 4))(x) # Reshape back to (16, 4)
x = keras.layers.Conv1DTranspose(8, 5, activation='linear', strides=2, padding='same')(x) # 16, 8
x = keras.layers.Conv1DTranspose(16, 5, activation='linear', strides=2, padding='same')(x) # 32, 16
x = keras.layers.Conv1DTranspose(8, 7, activation='linear', padding='same')(x) # 32, 16
decoded = keras.layers.Conv1DTranspose(1, 3, activation='linear', padding='same')(x) # 64, 1

decoder = keras.models.Model(encoded_input, decoded, name="decoder")

# Define the autoencoder
autoencoder_input = keras.layers.Input(shape=(64, 1), name="autoencoder_input")
encoded_ts = encoder(autoencoder_input)
decoded_ts = decoder(encoded_ts)

autoencoder = keras.models.Model(autoencoder_input, decoded_ts, name="autoencoder")

# Compile the autoencoder
autoencoder.compile(optimizer= Adam(learning_rate=0.000001), loss=weighted_mse_loss, metrics=[prd_loss_dig2phy_new])
weight_model = load_qmodel('inter_quantised.h5', custom_objects={'prd_loss_dig2phy_new': prd_loss_dig2phy_new, 'weighted_mse_loss': weighted_mse_loss})
# weight_model is a trained model almost identical to this model except for 
# the activation of the Dense layer at the end of the encoder being 'linear' instead of 'quantized_bits'
# Summary of the autoencoder
autoencoder.summary()
autoencoder.set_weights(weight_model.get_weights())
autoencoder.fit(features_train, new_train, epochs=30, shuffle=True, callbacks=[checkpoint_callback], validation_data=(features_val, new_val))