Package: onnx2torch
onnx2torch is an ONNX to PyTorch converter.
Our converter:Is easy to use – Convert the ONNX model with the function callconvert;Is easy to extend – Write your own custom layer in PyTorch and register it with@add_converter;Convert back to ONNX – You can convert the model back to ONNX using thetorch.onnx.exportfunction.If you find an issue, pleaselet us know! And feel free to create merge requests.Please note that this converter covers only a limited number of PyTorch / ONNX models and operations.
Let us know which models you use or want to convert from onnx to torchhere.Installationpipinstallonnx2torchorcondainstall-cconda-forgeonnx2torchUsageBelow you can find some examples of use.Convertimportonnximporttorchfromonnx2torchimportconvert# Path to ONNX modelonnx_model_path='/some/path/mobile_net_v2.onnx'# You can pass the path to the onnx model to convert it or...torch_model_1=convert(onnx_model_path)# Or you can load a regular onnx model and pass it to the converteronnx_model=onnx.load(onnx_model_path)torch_model_2=convert(onnx_model)ExecuteWe can execute the returnedPyTorch modelin the same way as the original torch model.importonnxruntimeasort# Create example datax=torch.ones((1,2,224,224)).cuda()out_torch=torch_model_1(x)ort_sess=ort.InferenceSession(onnx_model_path)outputs_ort=ort_sess.run(None,{'input':x.numpy()})# Check the Onnx output against PyTorchprint(torch.max(torch.abs(outputs_ort-out_torch.detach().numpy())))print(np.allclose(outputs_ort,out_torch.detach().numpy(),atol=1.e-7))ModelsWe have tested the following models:Segmentation models:DeepLabv3plusDeepLabv3 resnet50 (torchvision)HRNetUNet (torchvision)FCN resnet50 (torchvision)lraspp mobilenetv3 (torchvision)Detection  from MMdetection:SSDLite with MobileNetV2 backboneRetinaNet R50SSD300 with VGG backboneYolov3_d53Yolov5Classification fromtorchvision:Resnet18Resnet50MobileNet v2MobileNet v3 largeEfficientNet_b{0, 1, 2, 3}WideResNet50ResNext50VGG16GoogleleNetMnasNetRegNetTransformers:VitSwinGPT-J:page_facing_up: List of currently supported operations can be foundedhere.How to add new operations to converterHere we show how to extend onnx2torch with new ONNX operation, that supported by both PyTorch and ONNXand has the same behaviourAn example of such a module isRelu@add_converter(operation_type='Relu',version=6)@add_converter(operation_type='Relu',version=13)@add_converter(operation_type='Relu',version=14)def_(node:OnnxNode,graph:OnnxGraph)->OperationConverterResult:returnOperationConverterResult(torch_module=nn.ReLU(),onnx_mapping=onnx_mapping_from_node(node=node),)Here we have registered an operation namedRelufor opset versions 6, 13, 14.
Note that thetorch_moduleargument inOperationConverterResultmust be a torch.nn.Module, not just a callable object!
If Operation's behaviour differs from one opset version to another, you should implement it separately.but has different behaviourAn example of such a module isScatterND# It is recommended to use Enum for string ONNX attributes.classReductionOnnxAttr(Enum):NONE='none'ADD='add'MUL='mul'classOnnxScatterND(nn.Module,OnnxToTorchModuleWithCustomExport):def__init__(self,reduction:ReductionOnnxAttr):super().__init__()self._reduction=reduction# The following method should return ONNX attributes with their values as a dictionary.# The number of attributes, their names and values depend on opset version;# method should return correct set of attributes.# Note: add type-postfix for each key: reduction -> reduction_s, where s means "string".def_onnx_attrs(self,opset_version:int)->Dict[str,Any]:onnx_attrs:Dict[str,Any]={}# Here we handle opset versions < 16 where there is no "reduction" attribute.ifopset_version<16:ifself._reduction!=ReductionOnnxAttr.NONE:raiseValueError('ScatterND from opset < 16 does not support'f'reduction attribute !={ReductionOnnxAttr.NONE.value},'f'got{self._reduction.value}')returnonnx_attrsonnx_attrs['reduction_s']=self._reduction.valuereturnonnx_attrsdefforward(self,data:torch.Tensor,indices:torch.Tensor,updates:torch.Tensor,)->torch.Tensor:def_forward():# ScatterND forward implementation...returnoutputiftorch.onnx.is_in_onnx_export():# Please follow our convention, args consists of:# forward function, operation type, operation inputs, operation attributes.onnx_attrs=self._onnx_attrs(opset_version=get_onnx_version())returnDefaultExportToOnnx.export(_forward,'ScatterND',data,indices,updates,onnx_attrs)return_forward()@add_converter(operation_type='ScatterND',version=11)@add_converter(operation_type='ScatterND',version=13)@add_converter(operation_type='ScatterND',version=16)def_(node:OnnxNode,graph:OnnxGraph)->OperationConverterResult:node_attributes=node.attributesreduction=ReductionOnnxAttr(node_attributes.get('reduction','none'))returnOperationConverterResult(torch_module=OnnxScatterND(reduction=reduction),onnx_mapping=onnx_mapping_from_node(node=node),)Here we have used a trick to convert the model from torch back to ONNX by defining the custom_ScatterNDExportToOnnx.Opset version workaroundIncase you are using a model with older opset, try the following workaround:ONNX Version Conversion - Official DocsExampleimportonnxfromonnximportversion_converterimporttorchfromonnx2torchimportconvert# Load the ONNX model.model=onnx.load('model.onnx')# Convert the model to the target version.target_version=13converted_model=version_converter.convert_version(model,target_version)# Convert to torch.torch_model=convert(converted_model)torch.save(torch_model,'model.pt')Note: use this only when the model does not convert to PyTorch using the existing opset version. Result might vary.CitationTo cite onnx2torch useCite this repositorybutton, or:@misc{onnx2torch,
  title={onnx2torch},
  author={ENOT developers and Kalgin, Igor and Yanchenko, Arseny and Ivanov, Pyoter and Goncharenko, Alexander},
  year={2021},
  howpublished={\url{https://enot.ai/}},
  note={Version: x.y.z}
}AcknowledgmentsThanks to Dmitry Chudakov@cakeofwar42for his contributions.Special thanks to Andrey Denisov@denisovap2013for the logo design.