quic / aimet

AIMET is a library that provides advanced quantization and compression techniques for trained neural network models.
https://quic.github.io/aimet-pages/index.html
Other
2.08k stars 373 forks source link

Sample Code: Loading Custom Model/Checkpoint and Dataset for Post-Training Quantization (PTQ) #2591

Open adarshm93 opened 9 months ago

adarshm93 commented 9 months ago

Hello,

I am currently experimenting with the sample codes provided in the documentation/examples. I'm utilizing the default ResNet50 model, along with the dataset loader and evaluator functions. Currently, my objective is to load a TensorFlow checkpoint/model instead of relying on the pre-trained ResNet50 model weights from the Keras package.

However, While loading ckpt with "tf.saved_model.load()" I have encountered an issue: the "model.update" attribute is unavailable, preventing me from proceeding with further processes. As a result, I am unable to progress with the Post-Training Quantization (PTQ) process for a custom checkpoint. Despite my efforts to find documentation for this custom process, I have been unsuccessful in locating relevant content on the website. I kindly request any material or links that could guide me through the process as a reference. Thanks !!

quic-ernst commented 9 months ago

Hi @adarshm93 do you have a code snippet of your process?

adarshm93 commented 9 months ago

Appreciate your response @quic-ernst . I've included a code snippet where I load a saved resnet50 checkpoint (.pb) file and attempt to execute the quantization simulation process. However, I'm encountering challenges due to the lack of adequate examples for custom models and checkpoints. I would appreciate your suggestions and assistance on this matter.

Thank you !!

Screenshot (501)

quic-ernst commented 9 months ago

@adarshm93 I'm not too familiar with the MetaGraphDef that comes back from tf.saved_model.load call but I don't believe there is an update attribute. I've tried on my end and got the same result and I don't see any documentation about it.

Is your goal to load your checkpoint into a TF1 Session and try PTQ techniques with that model? If so, here is a Jupyter Notebook that walks through that process. Or is your model a Keras model and is a tf.keras.Model similar to the pretrained ResNet50 mentioned?

adarshm93 commented 9 months ago

@quic-ernst Yes, my objective is to load the checkpoint into a TF session and apply PTQ technique to that model. I followed the same notebook you mentioned, with the only modification being loading of the model (.pb) instead of fetching ResNet50 ImageNet weights from 'tf.keras.application' for my specific use case. Despite this adjustment, I am encountering the errors mentioned and struggling to find the appropriate information.

quic-ernst commented 9 months ago

@adarshm93 I've attached a code snippet of what I think your process should be. I don't think we need the update_ops_name that you had mentioned in the first post. We just need to load your session and continue tofold_all_batch_norms. This function will parse your session and fold all batch norms that it can. If it can't, it makes those non trainable and gives back the new session. From there, you can pass it to QuantizationSimModel and continue with the normal flow for AIMET.

import tensorflow.compat.v1 as tf

from aimet_tensorflow.batch_norm_fold import fold_all_batch_norms
from aimet_tensorflow.quantsim import QuantizationSimModel

sess = tf.Session()
loaded = tf.saved_model.load(export_dir="resnet50_saved_model", tags={"serve"}, sess=sess)

starting_op_names = # INPUT NAMES
output_op_names = # OUTPUT NAMES

sess, _ = fold_all_batch_norms(
    sess,
    input_op_names=starting_op_names,
    output_op_names=output_op_names
)

sim = QuantizationSimModel(
    sess,
    starting_op_names=starting_op_names,
    output_op_names=output_op_names
)

...