pytorch导出onnx时遇到不支持的算子怎么解决

在使用pytorch模型训练完成之后,我们现在使用的比较多的一种方法是将pytorch模型转成onnx格式的模型中间文件,然后再根据使用的硬件来生成具体硬件使用的深度学习模型,比如TensorRT。
在从pytorch模型转为onnx时,我们可能会遇到部分算子无法转换的问题,本篇注意记录下解决方法。

在导出onnx时,如果出现报错的算子,可以先在下面的链接中查找onnx算子是否支持
https://github.com/onnx/onnx/blob/main/docs/Operators.md

pytorch中有,onnx中也有的算子

导出时使用的onnx op 版本低导致

这个就好解决了,把op库的版本提高就行,但是有可能提高了版本以后,又出现了原来支持的算子现在又不支持了,这个再说

pytorch中没有注册某个onnx算子

如果是这种情况,就按照下面的方式进行:

from torch.onnx import register_custom_op_symbolic
# 创建一个asinh算子的symblic,符号函数,用来登记
# 符号函数内部调用g.op, 为onnx计算图添加Asinh算子
#   g: 就是graph,计算图
#   也就是说,在计算图中添加onnx算子
#   由于我们已经知道Asinh在onnx是有实现的,所以我们只要在g.op调用这个op的名字就好了
#   symblic的参数需要与Pytorch的asinh接口函数的参数对齐
#       def asinh(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...
def asinh_symbolic(g, input
### 回答1: PyTorch 中的 while 语句在导出 ONNX 需要使用特殊的函数进行处理,这个函数叫 torch.ops.script_ops.while_loop。该函数接受三个参数:循环条件、循环体和循环初始值。 示例: ```python import torch import torch.onnx def my_loop(counter, threshold): while counter < threshold: counter += 1 return counter counter = torch.tensor(0, dtype=torch.float32) threshold = torch.tensor(5, dtype=torch.float32) output = torch.ops.script_ops.while_loop( lambda counter, threshold: counter < threshold, my_loop, (counter, threshold) ) ``` 在上面的代码中,我们定义了一个 while 循环,并使用 torch.ops.script_ops.while_loop 函数将其换为 ONNX 可以使用的格式。 注意:这里的while语句内部只能使用torch支持的算子,否则会报错 ### 回答2: 在PyTorch中,通过将模型导出ONNX格式,while循环语句的处理略有不同。由于ONNX对于循环结构的支持相对有限,因此需要将while循环换为等效的数学表达式或通过迭代的方式来实现。 一种常见的处理方法是使用递归函数来替代while循环。首先,在导出将while循环替换为递归函数的调用。然后,在递归函数中,编写循环迭代的逻辑以模拟原有的while循环。最后,当达到退出条件,递归函数将返回结果。 另一种处理方法是换为数学表达式。通过分析while循环的逻辑,将其换为数学表达式来实现。然后,将该表达式嵌入到导出ONNX模型中。 需要注意的是,由于ONNX的限制,无法直接将while循环作为原子操作导出。因此,需要根据具体的循环逻辑进行相应的换处理。此外,ONNX还需要明确指定循环的最大迭代次数或合适的宽限范围,以确保换后的模型在运行的效果与原模型保持一致。 总之,在导出PyTorch模型至ONNX,需要根据循环的具体逻辑,将while循环换为递归函数或数学表达式,以保持模型的功能和性能。 ### 回答3: 在PyTorch中,导出ONNX要处理while语句,可以使用`torch.onnx.TracedModule`和`torch.onnx.export`函数来实现。 首先,需要使用`torch.onnx.TracedModule`将带有while语句的模型换为可跟踪的模型。这可以通过在模型前面添加`torch.jit.trace_module`来完成。例如,假设我们有一个带有while循环的模型`model`: ```python import torch import torch.nn as nn class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() def forward(self, x): i = 0 while i < 10: x = x + i i += 1 return x model = MyModel() ``` 然后,我们可以使用`torch.onnx.TracedModule`和`torch.jit.trace_module`来将带有while语句的模型换为可跟踪的模型: ```python traced_model = torch.onnx.TracedModule(torch.jit.trace_module(model)) ``` 接下来,可以使用`torch.onnx.export`函数将可跟踪的模型导出ONNX格式: ```python torch.onnx.export(traced_model, (input,), "model.onnx") ``` 在导出ONNX文件中,while循环会变成对应的ONNX运算符,并以递归方式执行。例如,在上面的示例中,导出ONNX图中将包含循环的等价操作和一个条件节点。 总结起来,将PyTorch模型导出ONNX处理while语句的步骤如下: 1. 使用`torch.onnx.TracedModule`将包含while循环的模型换为可跟踪的模型。 2. 使用`torch.onnx.export`将可跟踪的模型导出ONNX格式。 这样,导出ONNX文件中的while语句将通过等价的ONNX运算符来表示。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值