Hello!
I have an onnx file ( ANONYMOUS - Filemail ) that I try to convert to a .plan.
At first, when I tried to convert it, I had this error:
[TensorRT] ERROR: Network must have at least one output
I tried the solution explained here: https://devtalk.nvidia.com/default/topic/1044046/tensorrt/-tensorrt-error-network-must-have-at-least-one-output/
My build function if now:
import pretrainedmodels
import torch
import pretrainedmodels.utils as utils
from torch.nn import DataParallel, Sequential
from utils import ListImagesDataset, append
from torch.utils.data import DataLoader
import argparse
from tqdm import tqdm
import h5py
from torch.autograd import Variable
import torch.onnx
import torchvision
import tensorflow as tf
#import uff
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
# The Onnx path is used for Onnx models.
def build_engine_onnx(model_file, max_batch_size):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = 15 << 20
builder.max_batch_size = max_batch_size
# Load the Onnx model and parse it in order to populate the TensorRT network.
with open(model_file, 'rb') as model:
parser.parse(model.read())
print(network.get_layer(network.num_layers - 1).get_output(0).shape)
network.mark_output(network.get_layer(network.num_layers - 1).get_output(0))
return builder.build_cuda_engine(network)
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--model_name", default="resnet152Max", type=str)
parser.add_argument("--onnx_path", default="model", type=str)
parser.add_argument("--output_dir", required=True, type=str)
parser.add_argument("--num_workers", default=8, type=int)
parser.add_argument("--batch_size", required=True, type=int)
args = parser.parse_args()
with build_engine_onnx(args.onnx_path, args.batch_size) as engine:
with open(args.model_name+'.plan', 'wb') as f:
print('ok')
f.write(engine.serialize())
The code runs now, but I realized that the issue was that the parser does not parse every layers in my onnx model, I am missing the last one.
This means that the shape of my output is:
(2400, 7, 7)
Whereas I would like it to be:
(2400, 1, 1)
Any idea why?
Thank you in advance!