PyTorch 到 ONNX 导出，不支持 ATen 运算符，onnxruntime 挂起

Question

我想将基于 roberta-base 的语言模型导出为 ONNX 格式。该模型使用 ROBERTA 嵌入并执行文本分类任务。

from torch import nn
import torch.onnx
import onnx
import onnxruntime
import torch
import transformers

来自日志：

17: pytorch: 1.10.2+cu113
18: CUDA: False
21: device: cpu
26: onnxruntime: 1.10.0
27: onnx: 1.11.0

PyTorch 导出

batch_size = 3
model_input = {
    'input_ids': torch.empty(batch_size, 256, dtype=torch.int).random_(32000),
    'attention_mask': torch.empty(batch_size, 256, dtype=torch.int).random_(2),
    'seq_len':  torch.empty(batch_size, 1, dtype=torch.int).random_(256)
}
model_file_path = os.path.join("checkpoints", 'model.onnx')

torch.onnx.export(da_inference.model,               # model being run
                  model_input,                         # model input (or a tuple for multiple inputs)
                  model_file_path,   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=11,          # the ONNX version to export the model to
                  operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK,
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input_ids', 'attention_mask', 'seq_len'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input_ids': {0 : 'batch_size'},
                                'attention_mask': {0 : 'batch_size'},
                                'seq_len': {0 : 'batch_size'},
                                'output' : {0 : 'batch_size'}},
                 verbose=True)

我知道如果包含在模型架构中，从 ATen（C++11 的张量库）转换某些运算符可能会出现问题 PyTorch Model Export to ONNX Failed Due to ATen。

如果我设置参数 operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK 则导出成功，这意味着 'leave as is ATen operators if not supported in ONNX'.

PyTorch 导出函数给我以下警告：

Warning: Unsupported operator ATen. No schema registered for this operator.
Warning: Shape inference does not support models with experimental operators: ATen

看起来模型中唯一未转换为 ONNX 的 ATen 运算符位于层 LayerNorm.weight 和 LayerNorm.bias 内（我有几个这样的层）：

 %1266 : Float(3, 256, 768, strides=[196608, 768, 1], requires_grad=0, device=cpu) = 
onnx::ATen[cudnn_enable=1, eps=1.0000000000000001e-05, normalized_shape=[768], operator="layer_norm"]
(%1265, %model.utterance_rnn.base.encoder.layer.11.output.LayerNorm.weight,
 %model.utterance_rnn.base.encoder.layer.11.output.LayerNorm.bias)
# /opt/conda/lib/python3.9/site-packages/torch/nn/functional.py:2347:0

比模型检查通过 OK:

model = onnx.load(model_file_path)
# Check that the model is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))

我还可以使用 Netron 可视化计算图。

但是当我尝试使用导出的 ONNX 模型执行推理时，它停止了，没有日志或标准输出。所以这段代码会挂掉系统：

model_file_path = os.path.join("checkpoints", "model.onnx")
sess_options = onnxruntime.SessionOptions()
sess_options.log_severity_level = 0
ort_providers: List[str] = ["CUDAExecutionProvider"] if use_gpu else ['CPUExecutionProvider']
session = InferenceSession(model_file_path, providers=ort_providers, sess_options=sess_options)

有没有解决这个问题的建议？从官方文档中我看到 torch.onnx 以这种方式导出的模型可能只能由 Caffe2.

运行

这些层不在基础冷冻 roberta 模型中，所以这是我自己添加的附加层。是否可以用相似层替换有问题的层并重新训练模型？

或者 Caffe2 是这里的最佳选择并且 onnxruntime 不会进行推理？

更新：我在BERT cased embeddings的基础上重新训练了模型，但是问题依旧。相同的 ATen 运算符不会在 ONNX 中转换。看起来层 LayerNorm.weight 和 LayerNorm.bias 仅在 BERT 之上的模型中。那么，您对更改此层并启用 ONNX 导出有何建议？

Answer 1

您是否尝试过在为 onnx 定义运算符后导出？类似于以下内容 code by Huawei.

另一方面，加载模型时，从技术上讲，您可以覆盖任何您想要的内容。放置一个特定层等于您修改后的 class 继承原始层，保持相同的行为（输入和输出）但可以修改它的执行。您可以尝试使用它来保存更改后有问题的运算符的模型，在 onnx 中对其进行转换，并以这种形式（甚至在 pytorch 中）进行微调。

这通常似乎最好由 onnx 团队解决，因此长期解决方案可能是 post 在 github issues page 上对特定运算符的请求（但可能很慢）。

Answer 2

最好的方法是重写模型中使用这些运算符的地方，以某种方式转换查看 this 以供参考。例如，如果问题是层规范，那么您可以自己编写。有时有帮助的另一件事是不要将轴设置为动态的，因为有些操作还不支持它

PyTorch 到 ONNX 导出，不支持 ATen 运算符，onnxruntime 挂起

PyTorch to ONNX export, ATen operators not supported, onnxruntime hangs out

pytorch

caffe2

onnx