Implementing Post-training Quantization & Quantization-aware Training on a Pix2pix Image to Image Translation Model

Job ID: 37781576

Budget: $30 – $250 USD

I’m in need of implementing post-training quantization (Dynamic Range, Full Integer and Float16 Quantization), and quantization-aware training on a Pix2pix image to image translation source code written using TensorFlow 1.15. The aim of this is to reduce the model size.

This source code was implemented with Python 2.7 and as Google Colab doesn't support Python 2.7 by default, by creating a custom environment, I was able to run the code and get the results.

When the model is getting trained after every specific number of epochs I save a checkpoint and finally I restore the last checkpoint and produce test images. All these work fine but I need to implement quantization. So far I tried a method like below to convert the model to TFLite version:

===================================================
def convert(self, checkpoint_dir, checkpoint_path=None):
"""
Converts the TensorFlow model to TensorFlow Lite format.

Args:
- checkpoint_dir: Directory where checkpoints are saved.
- checkpoint_path: Optional specific checkpoint file to load. If None, the latest checkpoint in checkpoint_dir is used.
"""
# Ensure the latest checkpoint is loaded if checkpoint_path is not explicitly provided
if checkpoint_path is None:
checkpoint_path = tf.train.latest_checkpoint(checkpoint_dir)
if checkpoint_path is None:
print("No checkpoint found in", checkpoint_dir)
return
else:
checkpoint_path = os.path.join(checkpoint_dir, checkpoint_path)

# Attempt to restore the specified checkpoint
try:
self.saver.restore(self.sess, checkpoint_path)
print(f"Successfully loaded checkpoint: {checkpoint_path}")
except Exception as e:
print(f"Failed to load checkpoint from {checkpoint_path}. Error: {e}")
return

# Conversion process
try:
# Prepare the converter with the model's input and output
converter = tf.lite.TFLiteConverter.from_session(self.sess, [self.real_data], [self.d6])
# Set conversion parameters
converter.allow_custom_ops = True
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# Perform the conversion
tflite_model = converter.convert()

# Define the output path for the converted model
output_path = os.path.join(checkpoint_dir, "converted_model.tflite")

# Save the converted model to a .tflite file
with open(output_path, "wb") as f:
f.write(tflite_model)
print(f"Converted TFLite model written to {output_path}")
except Exception as e:
print(f"Failed to convert model. Error: {e}")
===================================================

Then another function like below to produce test results using the lite version of the model:

===================================================
def test_tflite(self, args):
"""Test the TFLite model."""
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="./results/checkpoints/converted_model.tflite")
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Retrieve and sort test image files
sample_files = glob("../datasets/train/*_yeniuydu_16bit_1.png")
sample_files = sorted(sample_files, key=lambda x: int(x.split('/')[-1].split('_')[0]))

for i, file_path in enumerate(sample_files):
print("Processing image:", file_path)

# Load and preprocess the image
test_image = load_data(file_path, is_test=True) # Assuming load_data is a standalone function
test_image = np.expand_dims(test_image, axis=0).astype(np.float32)

# Set input tensor
interpreter.set_tensor(input_details[0]['index'], test_image)

# Run the model
interpreter.invoke()

# Get the output
predictions = interpreter.get_tensor(output_details[0]['index'])

print(type(predictions))

Save the output images
for band in range(predictions.shape[-1] // 3):
start_idx = band * 3
end_idx = start_idx + 3
output_filename = f"{i+1}_3112020_test_ngf24l11000_{band+1}.png"
output_file_path = os.path.join(args.quantized_test_dir, output_filename)
save_images(
predictions[:, :, :, start_idx:end_idx],
[self.batch_size, 1],
output_file_path
)

print("TFLite testing finished.")
===================================================
The problem here is at all the images that I get are 2KB in size and all black (better to say the quantization) fails. I would like to work with a freelance who can implement the mentioned three types of the post-training quantization and quantization-aware training. To clarify these terms please read these descriptions:

Types of Quantization in TensorFlow Lite
Post-Training Quantization (PTQ): This is applied after a model has been trained. It reduces the precision of the weights and activations from floating-point to lower-bit representations, typically int8 or uint8. There are several types of post-training quantization:
Dynamic Range Quantization: The simplest form, where weights are quantized statically, but activations are quantized dynamically at runtime.
Full Integer Quantization: Converts both weights and activations to integers. This is useful for compatibility with integer-only hardware accelerators.
Float16 Quantization: Converts the weights to float16 instead of int8, reducing model size roughly by half with minimal loss in accuracy.
Quantization-Aware Training (QAT): This involves training the model with simulated quantization, meaning that the quantization effects are included in the forward and backward passes. This usually results in higher accuracy models compared to PTQ because the model can learn to adapt to the quantization induced noise.

Finally, you should pay attention that as the source is implemented using TensorFlow 1.15 some feature like "Converting a SavedModel to a TensorFlow Lite model" are not directly supported so you should whether implement them manually or migrate the code from TensorFlow 1.15 to 2.x.

I prefer to use the version 2.x of the TensorFlow but I don't want my source code to get totally changed. I'll share my source code with the freelancer.

For more information see this link:
https://www.tensorflow.org/api_docs/python/tf/lite/TFLiteConverter

Key requirements:
- Implement all three post-training quantization techniques.
- Implement quantization-aware training
- Optimize the model according to Google Colab platform guidelines.

Ideal skills for the project include:
- Proficiency in TensorFlow 1.15 or 2.x
- In-depth knowledge of post-training quantization methods.
- Expertise in model optimization for Google Colab.
- Knowledge in Pix2pix image to image translation with TensorFlow
Related categories: Python Image Processing Tensorflow