ONNX export in pytorch

[facebookresearch]

ONNX export in pytorch

[mnaumovfb]

Thank you for your comments. You are correct that the translation to ONNX has issues due to lack of ONNX support for some of the operators used to describe the model. Let me address some of your suggestions below.

1. There are tradeoffs associated with using Embedding and EmbeddingBag layers. Perhaps the most critical difference is that the inputs to the layers are defined differently. In particular, EmbeddingBag allows lookups with different number of indices to be easily batched together, while Embedding requires the number of indices in each lookup to be constant within a batch.

For the Kaggle Display Advertising Challenge Dataset this difference is irrelevant because each lookup has a single index in it, but the model is more general and can accept multiple indices per lookup (which can be controlled with a parameter from the command line). That is why we made a conscious choice to use the EmbeddingBag layer in the implementation.

2. This change seems reasonable. It make the code more compatible with ONNX at the expense of making it slightly more complicated. As you mentioned ultimately you will still hit a 2GB limit for buffer sizes for this dataset.

Therefore, if you are interested in saving protobuf without the parameters (weights/bias) then my advice is to try to use the Caffe2 version with an option "--save-proto-types-shapes", which should save the protobuf of the model including the shape and type of each of the operators. Alternatively, you can use the PyTorch version with an option "--save-model" and "--load-model" to save and load the model with parameters, respectively.