Hello,
I converted a standard resnet-18 pytorch model to onnx model using:
model = ResNet18(args).cuda()
print("pytorch to onnx")
# Translate Pytorch Model into Onnx Model
dummy_input = Variable(torch.randn(args.batch_size, args.input_channel, \
args.input_size, args.input_size, device='cuda'))
print("dummy input generated")
output_names = ["output"]
torch.onnx.export(model, dummy_input, args.onnx_model_name, verbose=False,
output_names=output_names)
Next, I checked the structure of the generated resnet18.onnx using:
# Load the ONNX model
model = onnx.load("resnet18.onnx")
# Check that the IR is well formed
onnx.checker.check_model(model)
# Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)
Noticed the information about the last a few layers:
%115[FLOAT, 512x512x3x3]
%116[FLOAT, 512]
%117[FLOAT, 512]
%118[FLOAT, 512]
%119[FLOAT, 512]
%120[INT64, scalar]
%121[FLOAT, 10x512]
%122[FLOAT, 10]
%188 = Padmode = ‘constant’, pads = [0, 0, 0, 0, 0, 0, 0, 0], value = 0
%189 = AveragePoolkernel_shape = [4, 4], pads = [0, 0, 0, 0], strides = [4, 4]
%190 = Constantvalue = <Scalar Tensor []>
%191 = Shape(%189)
%192 = Gather[axis = 0](%191, %190)
%193 = Constantvalue = <Scalar Tensor []>
%194 = Unsqueezeaxes = [0]
%195 = Unsqueezeaxes = [0]
%196 = Concat[axis = 0](%194, %195)
%197 = Reshape(%189, %196)
%output = Gemm[alpha = 1, beta = 1, transB = 1](%197, %121, %122)
return %output
However, when using tensorrt’s onnx parser to parse resnet18.onnx:
# The Onnx path is used for Onnx models.
def build_engine_onnx(model_file):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = common.GiB(1)
# 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)
I got the shape of (512, 1, 1) rather than (10, 1, 1) for the output layer, I also checked the shape of the second last layer which seems correct: (512, 4, 4). So looks like the onnx parser cannot parse the last layer correctly.
I found similar issue has been reported, but in my case this is a standard resnet18 model with input shape (3,32,32), can anyone help?
Thanks.