To enable quantization for a PyTorch model, it is necessary to modify the model definition to ensure it meets the following condition. Vitis AI GitHub provides an example.
- The model intended for quantization must contain only the forward method. All other functions, which typically serve as pre-processing and post-processing operations, should be moved outside the model or placed in a derived class. Failure to do so might result in unexpected behavior when forwarding the quantized module, as the API removes these functions in the quantized module.
- The float model should pass the JIT trace test. Set the
floating-point module to evaluation status, then use the
torch.jit.tracefunction to test the float model. For more details, refer to example/jupyter_notebook/jit_trace_test/jit_trace_test.ipynb.