import torch
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return torch.asinh(x)
from torch.onnx.symbolic_registry import register_op
def asinh_symbolic(g, input, *, out=None):
return g.op("Asinh", input)
register_op('asinh', asinh_symbolic, '', 10)
model = Model()
input = torch.rand(1, 3, 640, 640)
torch.onnx.export(model, input, 'asinh.onnx')
当运行如上代码,添加op算子后,把模型保存为onnx格式时,报错:
torch.onnx.symbolic_registry.UnsupportedOperatorError: Exporting the operator ::asinh to ONNX opset version 13 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.
解决办法:
register_op('asinh', asinh_symbolic, '', 10)
最后一个参数是opset版本,当设置为10时,需要在导出模型时设置相同的opset,即
torch.onnx.export(model, input, 'asinh.onnx', opset_version=10)