从 PyTorch 导出的 ONNX InferenceSession ONNX 模型失败

问题描述

我正在尝试将自定义 PyTorch 模型导出到 ONNX 以执行推理但没有成功...这里的棘手之处在于我正在尝试使用 基于脚本的导出器，如图所示示例 here 以便从我的模型中调用函数。

我可以毫无怨言地导出模型，但是在尝试启动 InferenceSession 时，我收到以下错误：

Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from ner.onnx Failed:Type Error: Type parameter (T) bound to different types (tensor(int64) and tensor(float) in node (Concat_1260).

我试图找出该问题的根本原因，它似乎是通过在以下函数中使用 torch.matmul() 生成的（非常讨厌，因为我试图仅使用 pytorch 运算符）：

@torch.jit.script
def valid_sequence_output(sequence_output,valid_mask):
    X = torch.where(valid_mask.unsqueeze(-1) == 1,sequence_output,torch.zeros_like(sequence_output))
    bs,max_len,_ = X.shape

    tu = torch.unique(torch.nonzero(X)[:,:2],dim=0)
    batch_axis = tu[:,0]
    rows_axis = tu[:,1]

    a = torch.arange(bs).repeat(batch_axis.shape).reshape(batch_axis.shape[0],-1)
    a = torch.transpose(a,1)

    T = torch.cumsum(torch.where(batch_axis == a,torch.ones_like(a),torch.zeros_like(a)),dim=1) - 1
    cols_axis = T[batch_axis,torch.arange(batch_axis.shape[0])]

    A = torch.zeros((bs,max_len))
    A[(batch_axis,cols_axis,rows_axis)] = 1.0

    valid_output = torch.matmul(A,X)
    valid_attention_mask = torch.where(valid_output[:,:,0] != 0,torch.ones_like(valid_mask),torch.zeros_like(valid_mask))
    return valid_output,valid_attention_mask

似乎不支持 torch.matmul（根据文档），所以我尝试了一系列解决方法（例如 A.matmul(X)、torch.baddbmm），但我仍然遇到同样的问题。 ..

有关如何解决此行为的任何建议都很棒：D 感谢您的帮助！

解决方法

这表明存在模型转换问题。请针对 Torch 导出器功能提出问题。类型 (T) 必须绑定到相同的类型才能使模型有效，而 ORT 基本上是在抱怨这一点。