PyTorch - 1.4.1 English

Vitis AI RNN User Guide (UG1563)

Document ID

UG1563

Release Date

2021-12-03

Version

1.4.1 English

The example/lstm_quant_pytorch/quantize_lstm.py file contains an example.

Import the PyTorch quantizer module.

from pytorch_nndct.apis import torch_quantizer

Generate a quantizer with quantization and get the converted model.

quantizer = torch_quantizer(quant_mode=args.quant_mode, 
                            module=model, 
                            bitwidth=16, 
                            lstm=True)
model = quantizer.quant_model

Forward a neural network with the converted model.
```
acc = test(model, DEVICE, test_loader)
```

Output the quantization result and deploy the model.

if args.quant_mode == 'calib':
    quantizer.export_quant_config()
if args.quant_mode == 'test':
    quantizer.export_xmodel(deploy_check=True)