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
orpytorch_pfn_extras.onnx.export
.- Parameters
fn (func) – the target function to be annotated,
args
is used for this function. Cannot passkwargs
for the function.args (tuple) – arguments for the target function
attrs (dict) – annotation paramters
- Return type
Any