pytorch_pfn_extras.onnx.annotate

pytorch_pfn_extras.onnx.annotate(**attrs)

Annotation parameters 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):
...         with pytorch_pfn_extras.onnx.annotate(key='value'):
...             h = self.conv(x)
...         h = self.conv2(h)
...         return h

Use this annotate function under with statement, then the first Conv operator will be emit with customized attributes. Customized attributes are invalid for ONNX format, so pay attention that some ONNX runtimes cannot run the output ONNX graph.

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

Parameters:

attrs (dict) – annotation parameters

Return type:

AbstractContextManager[None]