I am trying to run a pytorch neural network on the TX2 using TensorRT and I have been having problems at the stage of creating a tensorRT engine from the .onnx file.
For instance if I take the vgg16 network found here: https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py, and export it to the ONNX format like this:
import torch
import torchvision
mynetwork = torchvision.models.vgg.vgg16(pretrained=True)
input = torch.zeros((1, 3, 224, 224))
torch.onnx.export(mynetwork, input, "vgg16.onnx", verbose=True, input_names=["input"], output_names=["output"])
I get a file vgg16.onnx without any apparent problem. However when I feed it to TensorRT’s ONNX parser as is done in the sample, the program exits with the following message:
IndexError: Attribute not found: shape
It seems that this occurs during the call to
parser->convertToTRTNetwork()
Meanwhile, the sampleOnnxMNIST sample that came with tensorRT works just fine. Since the Readme for that sample reads as follows:
This sample demonstrates conversion of an MNIST network in ONNX format to
a TensorRT network. The network used in this sample can be found at https://github.com/onnx/models/tree/master/mnist
(model.onnx)
I went ahead and tried to run one of the model.onnx that is supplied there (there are 3, apparently for different ONNX opsets 1, 7 and 8), however none of them worked when used to create a tensorRT engine as in the sample. For instance, with the model.onnx with opset 7, I get the following:
[2019-02-21 02:03:24 ERROR] Parameter check failed at: ../builder/Network.cpp::addScale::112, condition: shift.count == 0 || shift.count == weightCount
python: onnx/converterToTRT.h:156: nvonnxparser::TRT_LayerOrWeights nvonnxparser::Converter::convert_node(const onnx::NodeProto*): Assertion `layer' failed.
Aborted (core dumped)
when calling parser->convertToTRTNetwork(). My guess is that the .onnx model supplied with the tensorRT samples is from previous versions of the ONNX models.
What ONNX formats / opsets are supported by tensorRT 4.0.2 on the TX2 and how can we obtain ONNX model that are compatible from pytorch models?
Some more information about my setup:
I installed tensorRT from the debian package tensorrt_4.0.2.0-1+cuda9.0_arm64.deb which I got from JetPack and I built and installed torch 1.1.0 from source on the TX2.