mirrors/mmclassification
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmclassification.git synced 2025-06-03 21:53:55 +08:00
Code Issues Projects Releases Wiki Activity
mmclassification/docs/en/tools/pytorch2torchscript.md
Ma Zerun 6f25bebe42
[Docs] Update docs. (#600) 
* [Docs] The template provides united Docs and OpenMMLab drop-lists in
banner. Remove them from our `conf.py`.

* Move docs of different languages to the same folder
2021-12-15 10:34:34 +08:00

1.7 KiB
Raw Blame History

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/pytorch2onnx.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