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

我正在使用Pytork来训练神经网络，并将它们输出到ONNX中。我在索引中使用这些模型，通过TensorRT加载NXS。我需要一个热编码的一些功能，但这是真的很难实现在Vespa框架

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

我已经注意到两件事：

• ONNX格式包括OneHot运算符：
• PyTorch内置ONNX导出系统不支持OneHot操作员：

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

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

---

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

或者，如果您可以从模型中提取转换，这样一个热编码张量就是网络的输入，那么您可以在Vespa端通过编写一个函数，通过将源数据转换为一个热张量来进行转换，例如

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

---

因此，根据我的测试，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()

该模型不做任何训练，只需要输入一个索引向量，一个hot使用PyTorch的

one_hot

对其进行编码，并将其发送到简单的NN层。权重是随机初始化的，我在这里的输出是：

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

这与Pytork在相同输入下给出的输出相同。因此PyTorch确实导出了ONNX操作符。这是针对PyTorch 1.7.1的

---

如果一个热编码的输入在Vespa中被索引为整数，那么您可以直接使用这些作为输入。

我可以用一些澄清来更好地帮助您。我不知道你说的通过TensorRT加载NX是什么意思。您是否使用外部模型服务器（TensorRT）并从Vespa查询结果发送功能？或者您正在使用Vespa的ONNX模型评估？无论如何，ONNX中的OneHot操作符只接受张量作为输入。您可以通过直接操作ONNX图（例如，使用python框架）来自行注入此功能。也许你可以解释一下这个用例？分类功能还有其他方式，例如PyTorch ONNX export支持的嵌入…谢谢@LesterSolbakken-我添加了一些关于我工作流程的详细信息。我的用例如下：Vespa计算一些浅层神经网络没有很好利用的分类功能，因此我想在训练之前，在这些功能的一个热编码版本（易于用Python编码）上训练我的模型。瓶颈在于Vespa对推理的定义（称为

searchDefinition

）并不能真正方便地再现通用的热编码。一个通过Python进行ONNX操作的例子对我真的很有帮助，你有吗？是的，我可以这样做，但对于生产使用来说，这并不是真正可持续的。我正在寻找一种更自动化的方式来包含OneHotencoding。请PyTorch处理此层也是一个不错的选择：谢谢！谢谢你，莱斯特，它很好用！请注意，您可以在ONNX中将分类输入转换为整数（这样您就可以保持searchDefinition不变）

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)]