Unverified Commit 192e29d5 authored by lin bin's avatar lin bin Committed by GitHub
Browse files

Add doc for quantizer export_model() (#3473)

parent 86d63a07
......@@ -110,12 +110,11 @@ Step2. Choose a quantizer and compress the model
Step3. Export compression result
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
You can export the quantized model directly by using ``torch.save`` api and the quantized model can be loaded by ``torch.load`` without any extra modification.
After training and calibration, you can export model weight to a file, and the generated calibration parameters to a file as well. Exporting onnx model is also supported.
.. code-block:: python
# Save quantized model which is generated by using NNI QAT algorithm
torch.save(model.state_dict(), "quantized_model.pth")
calibration_config = quantizer.export_model(model_path, calibration_path, onnx_path, input_shape, device)
Plese refer to :githublink:`mnist example <examples/model_compress/quantization/QAT_torch_quantizer.py>` for example code.
......
......@@ -92,6 +92,14 @@ def main():
train(model, quantizer, device, train_loader, optimizer)
test(model, device, test_loader)
model_path = "mnist_model.pth"
calibration_path = "mnist_calibration.pth"
onnx_path = "mnist_model.onnx"
input_shape = (1, 1, 28, 28)
device = torch.device("cuda")
calibration_config = quantizer.export_model(model_path, calibration_path, onnx_path, input_shape, device)
print("Generated calibration config is: ", calibration_config)
if __name__ == '__main__':
main()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment