pytorch_pfn_extras.onnx.apply_annotation#

pytorch_pfn_extras.onnx.apply_annotation(fn, *args, **attrs)#

Annotation applier to the target function

Usage:

>>> class Net(nn.Module):
...     def __init__(self):
...         super(Net, self).__init__()
...         self.conv = nn.Conv2d(1, 6, 3)
...         self.conv2 = nn.Conv2d(6, 12, 3)
...     def forward(self, x):
...         def _conv(x):
...             h = self.conv(x)
...             return torch.relu(h)
...         h = pytorch_pfn_extras.onnx.apply_annotation(
...             _conv, key='value')
...         h = self.conv2(h)
...         return h

Annotate into all operators emitted from the target function even if included not nn.Module function. On the above code, the first Conv and ReLu operator will be emit with customized attributes. Customized attributes are invlid for ONNX format, so pay attention that some ONNX runtimes cannot run the output ONNX graph.

This applier is enabled with either pytorch_pfn_extras.onnx.export_testcase or pytorch_pfn_extras.onnx.export.

Parameters:

  • fn (func) – the target function to be annotated, args is used for this function. Cannot pass kwargs for the function.

  • args (tuple) – arguments for the target function

  • attrs (dict) – annotation paramters

Return type:

Any