Pytorch通过保存为ONNX模型转TensorRT5的实现

Pytorch通过保存为通过保存为ONNX模型转模型转TensorRT5的实现的实现

主要介绍了Pytorch通过保存为ONNX模型转TensorRT5的实现，具有很好的参考价值，希望对大家有所帮助。

一起跟随小编过来看看吧

1 Pytorch以ONNX方式保存模型

def saveONNX(model, filepath):

'''

保存ONNX模型

:param model: 神经网络模型

2 利用TensorRT5中ONNX解析器构建Engine

def ONNX_build_engine(onnx_file_path):

'''

通过加载onnx文件，构建engine

:param onnx_file_path: onnx文件路径

:return: engine

'''

# 打印日志

G_LOGGER = trt.Logger(trt.Logger.WARNING)

with trt.Builder(G_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, G_LOGGER) as parser:

print('Completed parsing of ONNX file')

print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))

engine = builder.build_cuda_engine(network)

print("Completed creating Engine")

# 保存计划文件

# with open(engine_file_path, "wb") as f:

# f.write(engine.serialize())

return engine

3 构建TensorRT运行引擎进行预测

def loadONNX2TensorRT(filepath):

engine = self.ONNX_build_engine(filepath)

# 读取测试集

datas = DataLoaders()

test_loader = datas.testDataLoader()

img, target = next(iter(test_loader))

img = img.numpy()

target = target.numpy()

img = img.ravel()

context = engine.create_execution_context()

output = np.empty((100, 10), dtype=np.float32)