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

1 Pytorch以ONNX方式保存模型

def saveONNX(model, filepath):
 '''
 保存ONNX模型
 :param model: 神经网络模型
 :param filepath: 文件保存路径
 '''

# 神经网络输入数据类型
 dummy_input = torch.randn(self.config.BATCH_SIZE, 1, 28, 28, device='cuda')
 torch.onnx.export(model, dummy_input, filepath, verbose=True)

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:
  builder.max_batch_size = 100
  builder.max_workspace_size = 1 << 20

print('Loading ONNX file from path {}...'.format(onnx_file_path))
  with open(onnx_file_path, 'rb') as model:
   print('Beginning ONNX file parsing')
   parser.parse(model.read())
  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):
 '''
 通过onnx文件，构建TensorRT运行引擎
 :param filepath: onnx文件路径
 '''
 # 计算开始时间
 Start = time()

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)

# 分配内存
 d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
 d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
 bindings = [int(d_input), int(d_output)]

# pycuda操作缓冲区
 stream = cuda.Stream()
 # 将输入数据放入device
 cuda.memcpy_htod_async(d_input, img, stream)
 # 执行模型
 context.execute_async(100, bindings, stream.handle, None)
 # 将预测结果从从缓冲区取出
 cuda.memcpy_dtoh_async(output, d_output, stream)
 # 线程同步
 stream.synchronize()

print("Test Case: " + str(target))
 print("Prediction: " + str(np.argmax(output, axis=1)))
 print("tensorrt time:", time() - Start)

del context
 del engine

补充知识：Pytorch/Caffe可以先转换为ONNX，再转换为TensorRT

近来工作，试图把Pytorch用TensorRT运行。折腾了半天，没有完成。github中的转换代码，只能处理pytorch 0.2.0的功能（也明确表示不维护了）。和同事一起处理了很多例外，还是没有通过。吾以为，实际上即使勉强过了，能不能跑也是问题。

后来有高手建议，先转换为ONNX，再转换为TensorRT。这个思路基本可行。

是不是这样就万事大吉？当然不是，还是有严重问题要解决的。这只是个思路。

来源：https://blog.csdn.net/qq_38003892/article/details/89314108