pytorch导出onnx简单易懂(固定输入尺寸以及动态输入尺寸)

先简单定义一个网络结构

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.ClassifierLayer = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=4, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=4, out_channels=4, kernel_size=3, stride=1),
            nn.MaxPool2d(kernel_size=2),
        )


    def forward(self, x):
        y = self.ClassifierLayer(x)
        y = F.interpolate(y,size=x.shape[2:])
        return y

转换代码(固定输入尺寸)

#实例化网络结构
model = Net()
#定义一个模拟的输入
input_data = torch.randn([1,3,320,320])
#参数1:模型  参数2:模拟的输入数据  参数3:导出的onnx绝对路径  #参数4:是否打印导出过程中的详细信息
torch.onnx.export(model, input_data, "Net.onnx", verbose=True)

接下来可视化一下导出后的onnx网络结构
网络结构图
注意看右边的输入部分[1,3,320,320]

动态输入尺寸转换代码

    input_name = 'input'
    output_name = 'output'
    torch.onnx.export(model, 
                      input_data, 
                      "simplenet.onnx",
                      opset_version=11,
                      input_names=[input_name],
                      output_names=[output_name],
                      dynamic_axes={
                          input_name: {0: 'batch_size', 2: 'in_width', 3: 'int_height'},
                          output_name: {0: 'batch_size', 2: 'out_width', 3: 'out_height'}}
                      )

网络结构图
同样注意右边的输入部分,现在变成了[batch_size,3,in_width,int_height]

至此最基础的pytorch转onnx就全部完成了,所有的网络结构都可套用这两段代码,不过如果遇到了转换失败算子不支持的问题,可以尝试修改opset_version参数,或者使用其他结构去代替不支持的操作

再补充一下onnxsim:
如果你想把onnx再转其他的模型,就拿ncnn来举例,上图中的Gather在转换过程中就会出现不支持转到ncnn的错误信息。
推荐使用onnxsim,不论任何模型都可以用,当然能否sim成功还要看你的模型复杂程度,不过目前遇到的大部分模型都是可以经过sim规避很多问题的。
项目地址:onnxsim
假设你已经安装好了onnxsim,那么

onnxsim xxx.onnx xxx-sim.onnx

经过onnxsim之后,你的onnx模型就会简洁很多了,以静态输入尺寸举例如下图
请添加图片描述
可以看到ncnn不支持的Gather层已经被简化掉了。

注意:onnxsim的用法这里只说明了最基本的,其余的更多用法自行查看参数含义去使用。

  • 5
    点赞
  • 21
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论
### 回答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
发出的红包

打赏作者

木辛叔叔

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值