pytorch_pfn_extras.onnx.scoped_anchor

pytorch_pfn_extras.onnx.scoped_anchor(**attrs)

Add anchor node to the scoped modules

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.scoped_anchor(key='value'):
...             h = self.conv(x)
...         h = self.conv2(h)
...         return h

Use this scoped anchoring under with statement, then dummy Identity nodes are added before/after the first Conv operator with customized attributes.

This anchoring is triggered by nn.Module applying function, cannot use this with torch.* functions.

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]