Pytorch通过保存为ONNX模型转TensorRT5

最新推荐文章于 2025-04-16 21:18:18 发布

小关学长

最新推荐文章于 2025-04-16 21:18:18 发布

阅读量1.7w

点赞数 10

分类专栏：工作记录文章标签： TensorRT Pytorch ONNX

本文链接：https://blog.csdn.net/qq_38003892/article/details/89314108

版权

本文介绍了如何将Pytorch模型保存为ONNX格式，然后利用TensorRT5的ONNX解析器构建运行引擎，实现模型在TensorRT上的高效预测。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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_work