1 onnx导出模型出现RuntimeError: Exporting the operator pad_sequence to ONNX opset version 13 is not supported错误

今天在pytorch中导出模型为onnx时，由于使用了from torch.nn.utils.rnn import pad_sequence，在执行torch.onnx.export导出模型时出现了以下的错误

RuntimeError: Exporting the operator pad_sequence to ONNX opset version 13 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.

意思就是这个算法pad_seqeunce在onnx中不支持，我们可以自定义一下这个函数，就是将一个列表中不等长的Tensor自动补齐到最大长度，自定义的pad_seqeunce代码如下

def pad_sequence(sequences, batch_first=False, max_len=None, padding_value=0):
    """
        对一个List中的元素进行padding
            sequences:
            batch_first: 是否把batch_size放到第一个维度
            padding_value:
            max_len : 最大句子长度，默认为None，即在每个batch中以最长样本的长度对其它样本进行padding；
            当指定max_len的长度小于一个batch中某个样本的长度，那么在这个batch中还是会以最长样本的长度对其它样本进行padding
            建议指定max_len的值为整个数据集中最长样本的长度
        Returns:
        """
    max_size = sequences[0].size()
    trailing_dims = max_size[1:]
    length = max_len
    max_len = max([s.size(0) for s in sequences])
    if length is not None:
        max_len = max(length, max_len)
    if batch_first:
        out_dims = (len(sequences), max_len) + trailing_dims
    else:
        out_dims = (max_len, len(sequences)) + trailing_dims
    out_tensor = sequences[0].data.new(*out_dims).fill_(padding_value)
    for i, tensor in enumerate(sequences):
        length = tensor.size(0)
        # use index notation to prevent duplicate references to the tensor
        if batch_first:
            out_tensor[i, :length, ...] = tensor
        else:
            out_tensor[:length, i, ...] = tensor

    return out_tensor

使用自定义的pad_sequence函数替换掉from torch.nn.utils.rnn import pad_sequence函数，重新训练并导出为onnx即可。

参考连接