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

Question

我正在使用 PyTorch 训练神经网络并将它们输出到 ONNX。我在 Vespa 中使用这些模型index，它通过 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 确实支持单热编码导出到 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 的。

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