时间戳

首页 > 解决方案 > 将量化的 Pytorch 模型导出到 onnx 时出现分段错误

python - 将量化的 Pytorch 模型导出到 onnx 时出现分段错误

问题描述

我正在尝试将模型导出为该onnx格式。架构很复杂，所以我不会在这里分享，但基本上，我将网络权重放在一个.pth文件中。我能够加载它们，创建网络并使用它进行推理。需要注意的是，我已经调整了代码以能够量化网络。我添加了量化和去量化运算符以及一些torch.nn.quantized.FloatFunctional()运算符。

但是，每当我尝试使用

torch.onnx.export(torch_model,               # model being run
                  input_example,             # model input
                  model_name,                # where to save the model
                  export_params=True,        # store the trained parameter
                  opset_version=11,          # the ONNX version to export
                  # the model to
                  do_constant_folding=True,  # whether to execute constant
                  # folding for optimization
                 )

我知道Segmentation fault (core dumped) 我正在使用 Ubuntu 20，安装了以下软件包：

torch==1.6.0
torchvision==0.7.0
onnx==1.7.0
onnxruntime==1.4.0

请注意，根据prints我在代码中留下的一些内容，导出的推理部分完成。分段错误随后发生。

有谁知道为什么会发生这种情况？

[编辑]：当我的网络不适合量化操作时，我可以导出它。所以问题不是安装坏了，更多的是一些量化算子为了onnx节省的问题。

标签： pythonsegmentation-faultpytorchonnx

解决方案

好吧，事实证明 ONNX 不支持量化模型（但在运行时无论如何都不会警告你，它只是抛出一个段错误）。它似乎还没有提上日程，因此可以使用 TensorRT 来解决。

推荐阅读