# onnx2trt.pyimport tensorrt as trt
defONNX_build_engine(trt_model_name,onnx_model_name):
G_LOGGER = trt.Logger()
explicit_batch =1<<(int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)#trt7with trt.Builder(G_LOGGER)as builder, builder.create_network(explicit_batch)as network, trt.OnnxParser(network, G_LOGGER)as parser:
builder.max_batch_size =1
builder.max_workspace_size =1<<30print('Loading ONNX file from path {}...'.format(onnx_model_name))withopen(onnx_model_name,'rb')as model:print('Beginning ONNX file parsing')#b = parser.parse(model.read())ifnot parser.parse(model.read()):for error inrange(parser.num_errors):print(parser.get_error(error))if1:print('Completed parsing of ONNX file')print('Building an engine from file {}; this may take a while...'.format(onnx_model_name))##### builder.int8_mode = True# builder.int8_calibrator = calib
builder.fp16_mode =True####print("num layers:",network.num_layers)
last_layer = network.get_layer(network.num_layers -1)# if not last_layer.get_output(0):# network.mark_output(network.get_layer(network.num_layers - 1).get_output(0))# 有的模型需要,有的模型在转onnx的之后已经指定了,就不需要这行
network.get_input(0).shape =[1,3,640,640]# trt7
engine = builder.build_cuda_engine(network)print("engine:",engine)print("Completed creating Engine")withopen(trt_model_name,"wb")as f:
f.write(engine.serialize())return engine
else:print('Number of errors: {}'.format(parser.num_errors))
error = parser.get_error(0)# if it gets mnore than one error this have to be changeddel parser
desc = error.desc()
line = error.line()
code = error.code()print('Description of the error: {}'.format(desc))print('Line where the error occurred: {}'.format(line))print('Error code: {}'.format(code))print("Model was not parsed successfully")
exit(0)
ONNX_build_engine('engine路径','onnx路径')
运行命令:
python3 onnx2trt.py
[TensorRT] ERROR: Network must have at least one output
[TensorRT] ERROR: Network validation failed.
engine: None
Completed creating Engine
Traceback (most recent call last):
File "onnx2trt.py", line 49, in<module>
ONNX_build_engine('engine','onnx')
File "onnx2trt.py", line 33, in ONNX_build_engine
f.write(engine.serialize())
AttributeError: 'NoneType' object has no attribute 'serialize'
若出现,以上情况,说明输入输出的定位出现了问题,此时将代码中,
if not last_layer.get_output(0)这两行放开,则问题解决:
Loading ONNX file from path **.onnx...
Beginning ONNX file parsing
[TensorRT] WARNING: onnx2trt_utils.cpp:220: Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
Completed parsing of ONNX file
Building an engine from file **.onnx; this may take a while...
num layers: 368
engine: <tensorrt.tensorrt.ICudaEngine object at 0x7f65cef618>
Completed creating Engine