如何在来自 PyTorch 的 ONNX 中包含 OneHot

Question

我正在使用 PyTorch 训练神经网络并将它们输出到 ONNX。我在 Vespa 索引中使用这些模型，它通过 TensorRT 加载 ONNX。我需要对某些功能进行单热编码，但这在 Vespa 框架内真的很难实现。

是否可以为我的 ONNX 网络中的某些给定功能嵌入单热编码（例如，在网络表示之前）？如果是这样，我应该如何基于 PyTorch 模型实现这一点？

我已经注意到两件事：

ONNX 格式包括 OneHot 运算符：see ONNX doc
PyTorch 内置 ONNX 导出系统不支持 OneHot 算子：see torch.onnx doc

编辑 2021/03/11：这是我的工作流程：

通过 PyTorch 训练排序学习模型
将它们导出为 ONNX
将这些 ONNX 导入我的 Vespa 索引，以便借助 ONNX 模型对任何查询的结果进行排名。在引擎盖下，Vespa 使用 TensorRT 进行推理（所以我使用 Vespa 的 ONNX 模型评估）

Answer 1

如果 PyTorch 无法将 OneHot 运算符导出到 ONNX，我认为你最好的选择是让他们解决这个问题？

或者，如果你可以从你的模型中提取转换，这样 one-hot-encoded 张量就是你网络的输入，你可以在 Vespa 端通过编写一个提供 one-通过将源数据转换为热张量，例如

function oneHotInput() {
    expression: tensor(x[10])(x == attribute(myInteger))
}

Answer 2

因此，根据我的测试，PyTorch 确实支持单热编码导出到 ONNX。使用以下型号：

#! /usr/bin/env python3

import torch
import torch.onnx
import torch.nn.functional as F


class MyModel(torch.nn.Module):
    def __init__(self, classes=5):
        super(MyModel, self).__init__()
        self._classes = classes
        self.linear = torch.nn.Linear(in_features=self._classes, out_features=1)
        self.logistic = torch.nn.Sigmoid()

    def forward(self, input):
        one_hot = F.one_hot(input, num_classes=self._classes).float()
        return self.logistic(self.linear(one_hot))


def main():
    model = MyModel()

    # training omitted

    data = torch.tensor([0, 4, 2])
    torch.onnx.export(model, data, "test.onnx", 
        input_names=["input"], output_names=["output"])
        
    result = model.forward(data)  
    print(result)

if __name__ == "__main__":
    main()

该模型不进行任何训练，仅采用索引向量，使用 PyTorch 的 one_hot 对它们进行单热编码，并将其发送到简单的 NN 层。权重是随机初始化的，这里的输出对我来说是：

tensor([[0.5749],
        [0.5081],
        [0.5581]], grad_fn=<SigmoidBackward>)

此模型导出到 ONNX 到“test.onnx”文件。使用 ONNX Runtime（这是 Vespa 在后端使用的，而不是 TensorRT）测试该模型：

In [1]: import onnxruntime as ort                                                                                                                                                            
In [2]: m = ort.InferenceSession("test.onnx")                                                                                                                                                
In [3]: m.run(input_feed={"input":[0,4,2]}, output_names=["output"])                                                                                                                        
Out[3]: 
[array([[0.57486993],
        [0.5081395 ],
        [0.5580716 ]], dtype=float32)]

这与具有相同输入的 PyTorch 给出的输出相同。所以 PyTorch 确实导出了 OneHot ONNX 运算符。这是针对 PyTorch 1.7.1 的。

如果 one-hot 编码的输入在 Vespa 中被索引为整数，那么您可以直接将它们用作输入。

Answer 3

您可以在导出模型时选择设置 optset=7 或找到另一个

如何在来自 PyTorch 的 ONNX 中包含 OneHot

How to include a OneHot in an ONNX coming from PyTorch

one-hot-encoding

pytorch

vespa

onnx