Note

You are reading the documentation for MMClassification 0.x, which will soon be deprecated at the end of 2022. We recommend you upgrade to MMClassification 1.0 to enjoy fruitful new features and better performance brought by OpenMMLab 2.0. Check the installation tutorial, migration tutorial and changelog for more details.

Pytorch to TorchScript (Experimental)¶

Pytorch to TorchScript (Experimental)

How to convert models from Pytorch to TorchScript¶

Usage¶

python tools/deployment/pytorch2torchscript.py \
    ${CONFIG_FILE} \
    --checkpoint ${CHECKPOINT_FILE} \
    --output-file ${OUTPUT_FILE} \
    --shape ${IMAGE_SHAPE} \
    --verify \

Description of all arguments¶

config : The path of a model config file.
--checkpoint : The path of a model checkpoint file.
--output-file: The path of output TorchScript model. If not specified, it will be set to tmp.pt.
--shape: The height and width of input tensor to the model. If not specified, it will be set to 224 224.
--verify: Determines whether to verify the correctness of an exported model. If not specified, it will be set to False.

Example:

python tools/deployment/pytorch2torchscript.py \
    configs/resnet/resnet18_8xb16_cifar10.py \
    --checkpoint checkpoints/resnet/resnet18_8xb16_cifar10.pth \
    --output-file checkpoints/resnet/resnet18_8xb16_cifar10.pt \
    --verify \

Notes:

All models are tested with Pytorch==1.8.1

Reminders¶

For torch.jit.is_tracing() is only supported after v1.6. For users with pytorch v1.3-v1.5, we suggest early returning tensors manually.
If you meet any problem with the models in this repo, please create an issue and it would be taken care of soon.

FAQs¶

None