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)

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐