3

I'm using PyTorch to train neural-net and output them into ONNX. I use these models in a Vespa index, which loads ONNXs through TensorRT. I need one-hot-encoding for some features but this is really hard to achieve within the Vespa framework.

Is it possible to embed a one-hot-encoding for some given features inside my ONNX net (e.g. before the network's representation) ? If so, how should I achieve this based on a PyTorch model ?

I already noticed two things:

EDIT 2021/03/11: Here is my workflow:

  • training learning-to-rank models via PyTorch
  • exporting them as ONNX
  • importing these ONNX into my Vespa index in order to rank any query's results thanks to the ONNX model. Under the hood, Vespa uses TensorRT for inference (so I use Vespa's ONNX model evaluation)
fweber
  • 355
  • 1
  • 2
  • 10
  • 1
    I could use some clarification to better help you. I'm not sure what you mean by loading ONNX through TensorRT here. Are you using an external model server (TensorRT) and sending features from the results of a Vespa query? Or are you using Vespa's ONNX model evaluation? Anyway, the OneHot operator in ONNX just takes a tensor as input. It is possible to inject this yourself by manipulating the ONNX graph directly, using e.g. the python framework. Perhaps you could explain the use case? There are other ways for categorical features, such as embeddings supported by PyTorch ONNX export... – Lester Solbakken Mar 10 '21 at 08:23
  • Thanks @LesterSolbakken - I added some details about my workflow. My use case is the following: Vespa computes some categorical features that are not well exploited by shallow neural-nets, therefore I would like to train my models on one-hot-encoded versions of these features (easy to encode them in Python) prior to my training. The bottleneck is that Vespa's definition of the inference (called `searchDefinition`) is not really handy to reproduce generic one-hot-encoding. An example of ONNX manipulation through Python would really help me, do you have any ? – fweber Mar 11 '21 at 10:32

3 Answers3

2

So, according to my testing, PyTorch does support one-hot encoding export to ONNX. With the following model:

#! /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()

This model doesn't do any training, just takes a vector of indices in, one-hot encodes them using PyTorch's one_hot and sends that to the simple NN layer. The weights are randomly initialised, and the output here for me was:

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

This model is exported to ONNX to the "test.onnx" file. Testing this model using ONNX Runtime (which is what Vespa uses in the backend, not 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)]

Which is the same output as given from PyTorch with the same input. So PyTorch does export the OneHot ONNX operator. This was for PyTorch 1.7.1.

If the input to the one-hot encoding is indexed in Vespa as integers, you can then just use these directly as inputs.

Lester Solbakken
  • 266
  • 1
  • 2
  • Thanks Lester, it's working fine ! Note that you can cast the categorical input as integer *within* the ONNX (so you can leave your searchDefinition untouched) – fweber May 04 '21 at 10:38
0

If PyTorch can't export the OneHot operator to ONNX I think your best option is to ask them to fix that?

Or, if you can extract the conversion from your model, such that the one-hot-encoded tensor is an input to your network, you can do that conversion on the Vespa side by writing a function supplying the one-hot tensor by converting the source data to it, e.g

function oneHotInput() {
    expression: tensor(x[10])(x == attribute(myInteger))
}
Jon
  • 2,043
  • 11
  • 9
  • Yes I could do that but it's not really sustainable for a production use. I'm looking for a more automated way of including OneHotEncodings. Asking PyTorch to handle this layer is also a good option: thanks ! – fweber Mar 11 '21 at 10:34
0

you can chose setting optset=7 when you export model or find anot

胖头鱼
  • 1