今天在进行onnx的模型导出时,遇到了一个问题,如标题,就是“torch.roll"这个算子在从pytorch到onnx的转换过程中会失效。
对应的解决方法是,把pytorch的版本升级,原先的版本是1.8,后来升级到了1.10。然后解决了该问题。
但是引起我注意的是,如果使用旧版本的pytorch,还有一种解决方式,那就是基于该文件解决https://github.com/pytorch/pytorch/blob/main/torch/onnx/symbolic_opset9.py
其中,对应的内容如下
#This decorator marks the function as the symbolic implementation
#for the aten::roll operation in PyTorch,
#allowing it to be translated into the corresponding ONNX operation.
@_onnx_symbolic("aten::roll")
@symbolic_helper.parse_args("v", "is", "is")
#This decorator handles the parsing of function arguments.
#"v" means that the first argument (self) is a tensor value.
#"is" indicates that the second (shifts) and third (dims) arguments are lists of integers.
def roll(g: jit_utils.GraphContext, self, shifts, dims):
assert len(shifts) == len(dims)
result = self
for i in range(len(shifts)):
shapes = []
shape = symbolic_helper._slice_helper(
g, result, axes=[dims[i]], starts=[-shifts[i]], ends=[sys.maxsize]
)
shapes.append(shape)
shape = symbolic_helper._slice_helper(
g, result, axes=[dims[i]], starts=[0], ends=[-shifts[i]]
)
shapes.append(shape)
result = g.op("Concat", *shapes, axis_i=dims[i])
return result
关于这个过程,我的问题就是,既然torch中有torch.roll的实现, 为什么这里转换不行呢?
从该文档中我们可以寻找到答案。
在torch模型转换为onnx的过程中,使用的是torch.jit,ScrptModule,而不是torch.nn.Module,关于这两者的区别,有不少更好的视频介绍,简单理解为前者用于提供模型静态图推理,强调性能,后者强调“简洁,美丽”,是一种不关注性能的动态图。
在torch到onnx的转换过程中,torchscript图中的每一个节点会被exporter按照拓扑结构导出,每一个ops会被对应为一个registered symbolic function.
如果一个ops没有被定义,你需要在torch/onnx/symbolic_opset<version>.py中去提供定义。
关于进一步的使用介绍,建议参考TorchScript-based ONNX Exporter — PyTorch 2.4 documentation