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

【问题标题】：How to include a OneHot in an ONNX coming from PyTorch如何在来自 PyTorch 的 ONNX 中包含 OneHot
【发布时间】：2021-03-09 10:20:58
【问题描述】：

我正在使用 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 模型评估）

【问题讨论】：

我可以使用一些澄清来更好地帮助你。我不确定你在这里通过 TensorRT 加载 ONNX 是什么意思。您是否使用外部模型服务器 (TensorRT) 并从 Vespa 查询的结果中发送特征？或者您正在使用 Vespa 的 ONNX 模型评估？无论如何，ONNX 中的 OneHot 运算符只是将张量作为输入。可以通过直接操作 ONNX 图来自己注入，例如使用蟒蛇框架。也许您可以解释用例？分类特征还有其他方法，例如 PyTorch ONNX 导出支持的嵌入...
谢谢@LesterSolbakken - 我添加了一些关于我的工作流程的细节。我的用例如下： Vespa 计算一些分类特征，浅层神经网络没有很好地利用这些特征，因此我想在这些特征的单热编码版本上训练我的模型（易于在 Python 中编码）到我的训练。瓶颈在于 Vespa 的推理定义（称为searchDefinition）对于重现通用的单热编码并不是很方便。通过 Python 操作 ONNX 的示例对我很有帮助，你有吗？

标签： pytorch one-hot-encoding onnx vespa

【解决方案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 的。

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

【讨论】：

谢谢莱斯特，一切正常！请注意，您可以将分类输入转换为 ONNX within 的整数（因此您可以保持 searchDefinition 不变）

【解决方案2】：

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

或者，如果您可以从模型中提取转换，从而将 one-hot-encoded 张量作为网络的输入，您可以通过编写提供 one-hot 张量的函数在 Vespa 端进行转换通过将源数据转换为它，例如

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

【讨论】：

是的，我可以这样做，但它对于生产用途来说并不真正可持续。我正在寻找一种更自动化的方式来包含 OneHotEncodings。让 PyTorch 处理这一层也是一个不错的选择：谢谢！

【解决方案3】：

您可以在导出模型时选择设置 optset=7 或者找不到

【讨论】：