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 withtorch.*
functions.This annotation is enabled with either
pytorch_pfn_extras.onnx.export_testcase
orpytorch_pfn_extras.onnx.export
.- Parameters
attrs (dict) – annotation parameters