onnx2torch2 2.2.0

An example of such a module is Relu

@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:
    return OperationConverterResult(
        torch_module=nn.ReLU(),
        onnx_mapping=onnx_mapping_from_node(node=node),
    )

Here we have registered an operation named Relu for opset versions 6, 13, 14. Note that the torch_module argument in OperationConverterResult must 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.

An example of such a module is ScatterND

# It is recommended to use Enum for string ONNX attributes.
class ReductionOnnxAttr(Enum):
    NONE = "none"
    ADD = "add"
    MUL = "mul"


class OnnxScatterND(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.
        if opset_version < 16:
            if self._reduction != ReductionOnnxAttr.NONE:
                raise ValueError(
                    "ScatterND from opset < 16 does not support"
                    f"reduction attribute != {ReductionOnnxAttr.NONE.value},"
                    f"got {self._reduction.value}"
                )
            return onnx_attrs

        onnx_attrs["reduction_s"] = self._reduction.value
        return onnx_attrs

    def forward(
        self,
        data: torch.Tensor,
        indices: torch.Tensor,
        updates: torch.Tensor,
    ) -> torch.Tensor:
        def _forward():
            # ScatterND forward implementation...
            return output

        if torch.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())
            return DefaultExportToOnnx.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.attributes
    reduction = ReductionOnnxAttr(node_attributes.get("reduction", "none"))
    return OperationConverterResult(
        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.

import onnx
from onnx import version_converter
import torch
from onnx2torch2 import convert

# Load the ONNX model.
model = onnx.load("model.onnx")
# Convert the model to the target version.
target_version = 13
converted_model = version_converter.convert_version(model, target_version)
# Convert to torch.
torch_model = convert(converted_model)
torch.save(torch_model, "model.pt")