DeOldify跨框架模型转换：从PyTorch到ONNX及TensorRT加速

DeOldify跨框架模型转换从PyTorch到ONNX及TensorRT加速最近在折腾一个挺有意思的项目想把老照片上色的模型DeOldify部署到生产环境里。原版模型是用PyTorch写的直接拿来用的话推理速度总觉得差点意思尤其是在处理大批量图片的时候。相信不少朋友也遇到过类似的问题训练好的模型怎么才能让它跑得更快、更稳我花了一些时间研究发现走PyTorch → ONNX → TensorRT这条技术路线效果相当不错。整个过程就像给模型做了一次“深度优化手术”不仅让它能在不同框架间自由穿梭还能借助TensorRT的加速能力让推理速度提升好几倍。今天我就把这个从模型导出到加速部署的完整过程一步步拆开讲给你听。即使你之前没怎么接触过模型转换跟着做下来也能搞定。1. 为什么需要模型转换与加速在开始动手之前我们先聊聊为什么要费这个劲。你可能会想PyTorch模型用得好好的干嘛要转来转去简单来说这主要是为了解决两个核心问题部署灵活性和推理性能。PyTorch非常适合研究和快速实验但到了实际部署尤其是对延迟和吞吐量要求高的场景比如在线服务、移动端应用它可能就不是最优选了。ONNXOpen Neural Network Exchange就像一个“中间翻译官”它定义了一个通用的模型表示格式。把PyTorch模型转成ONNX就意味着这个模型可以在支持ONNX的众多推理引擎上运行比如TensorRT、OpenVINO、ONNX Runtime等不再被PyTorch“绑定”。而TensorRT是NVIDIA推出的高性能深度学习推理SDK。它会对模型进行一系列深度的优化包括层融合、精度校准如FP16/INT8、内核自动调优等并为特定的GPU硬件生成高度优化的推理引擎。经过TensorRT优化后的模型推理速度通常能有数倍甚至数十倍的提升。所以对于我们的DeOldify上色模型这条路径的价值就很明确了先通过ONNX获得框架无关的便携性再通过TensorRT榨取GPU的极限性能最终实现高效、稳定的生产级部署。2. 准备工作与环境搭建工欲善其事必先利其器。我们先来把需要的工具和环境准备好。2.1 模型与代码获取首先你需要有训练好的DeOldify模型权重.pth文件以及对应的模型定义代码。你可以从DeOldify的官方仓库获取源码和预训练模型。这里假设你已经有了一个可用的PyTorch模型并且知道如何加载它进行推理。2.2 核心工具包安装我们将主要用到以下几个Python库PyTorch: 用于加载原始模型和进行ONNX导出。onnx / onnxruntime: ONNX格式的支持库和运行时。torch.onnx: PyTorch自带的ONNX导出模块。onnx-simplifier: 用于简化导出的ONNX模型去除冗余操作对后续转换至关重要。TensorRT: NVIDIA的推理加速库。安装它稍微复杂一些。我强烈建议在一个干净的Python虚拟环境中进行以下操作。这里以使用pip安装为例# 安装PyTorch (请根据你的CUDA版本选择合适命令这里以CUDA 11.8为例) pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 安装ONNX相关库 pip install onnx onnxruntime onnx-simplifier # 安装TensorRT Python包。 # 注意通常需要先从NVIDIA官网下载TensorRT的.tar.gz安装包解压后找到python目录下的whl文件进行安装。 # 例如pip install TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp3x-none-linux_x86_64.whl # 同时需要将TensorRT的lib目录添加到LD_LIBRARY_PATH环境变量。2.3 验证环境安装完成后可以写个简单的脚本验证一下关键库是否能正常导入import torch import onnx import onnxsim import tensorrt as trt print(fPyTorch version: {torch.__version__}) print(fONNX version: {onnx.__version__}) print(fTensorRT version: {trt.__version__})如果以上导入都没有报错那么恭喜你基础环境就准备好了。接下来我们进入核心的转换流程。3. 第一步将PyTorch模型导出为ONNX这是整个流程的起点。目标是将动态的PyTorch模型“冻结”成一个静态计算图并保存为ONNX格式。3.1 模型加载与准备假设你的模型定义在一个叫DeOldify的类中。首先你需要像平常一样加载模型权重并将模型设置为评估模式。这一点很重要因为像Dropout、BatchNorm这样的层在推理和训练时的行为是不同的。import torch from your_model_definition import DeOldify # 替换为你的模型定义 # 1. 初始化模型结构 model DeOldify(...) # 填入你的模型初始化参数 # 2. 加载预训练权重 checkpoint torch.load(your_deoldify_model.pth, map_locationcpu) # 注意权重文件的key可能需要根据你的保存方式调整例如 model.load_state_dict(checkpoint[state_dict]) model.load_state_dict(checkpoint) # 3. 切换到评估模式 model.eval() # 4. 将模型移动到GPU如果可用 device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device)3.2 构造示例输入与动态轴设置ONNX导出需要提供一个示例输入dummy inputPyTorch会根据这个输入进行一次前向传播来追踪计算图。对于图像模型输入通常是四维张量[批次大小, 通道数, 高度, 宽度]。为了让导出的模型能支持可变的批次大小或图像尺寸这在部署中很常见我们需要设置动态轴。# 构造一个示例输入张量 batch_size 1 # 导出时可以用1通过动态轴支持变化 height, width 256, 256 # 根据你的模型输入尺寸调整 dummy_input torch.randn(batch_size, 3, height, width).to(device) # 定义输入/输出的动态轴。 # 这里将输入和输出张量的第0维批次维设置为动态。 dynamic_axes { input: {0: batch_size}, # 输入名: {维度: ‘名字’} output: {0: batch_size} # 输出名: {维度: ‘名字’} } # 注意input和output需要与模型forward方法的参数名和返回的字典key如果有对应。3.3 执行ONNX导出使用torch.onnx.export函数进行导出。你需要指定模型、示例输入、导出路径、输入输出名称、动态轴配置以及操作集版本。# 导出ONNX模型 onnx_model_path deoldify.onnx torch.onnx.export( model, # 要导出的模型 dummy_input, # 模型输入或元组 onnx_model_path, # 保存路径 input_names[input], # 输入节点名称 output_names[output], # 输出节点名称 dynamic_axesdynamic_axes, # 动态轴配置 opset_version14, # ONNX操作集版本建议11 do_constant_foldingTrue, # 是否进行常量折叠优化 verboseFalse # 是否打印详细信息 ) print(f模型已导出至: {onnx_model_path})导出完成后建议使用onnx库加载一下检查模型结构是否正常并验证一下输出。import onnx # 加载并检查模型 onnx_model onnx.load(onnx_model_path) onnx.checker.check_model(onnx_model) print(ONNX模型检查通过) # (可选) 使用ONNX Runtime进行简单推理验证 import onnxruntime as ort import numpy as np ort_session ort.InferenceSession(onnx_model_path) ort_inputs {ort_session.get_inputs()[0].name: dummy_input.cpu().numpy()} ort_outputs ort_session.run(None, ort_inputs) print(ONNX Runtime推理成功输出shape:, ort_outputs[0].shape)4. 第二步优化与简化ONNX模型直接从PyTorch导出的ONNX模型可能包含一些冗余的操作或节点这会影响后续TensorRT转换的成功率和效率。onnx-simplifier工具可以自动简化模型图结构。import onnx from onnxsim import simplify # 加载原始ONNX模型 model onnx.load(onnx_model_path) # 简化模型 # dynamic_input_shape 选项告诉简化器输入形状是动态的 model_simp, check simplify(model, dynamic_input_shapeTrue) if check: simplified_path deoldify_simplified.onnx onnx.save(model_simp, simplified_path) print(f模型简化成功保存至: {simplified_path}) # 计算简化前后模型大小 import os orig_size os.path.getsize(onnx_model_path) / 1024**2 simp_size os.path.getsize(simplified_path) / 1024**2 print(f原始模型大小: {orig_size:.2f} MB) print(f简化后模型大小: {simp_size:.2f} MB) else: print(模型简化失败) simplified_path onnx_model_path # 如果简化失败使用原始模型简化后的模型通常更干净转换到TensorRT时遇到的兼容性问题会更少。强烈建议将此步骤作为标准流程。5. 第三步使用TensorRT构建加速引擎现在我们有了一个“干净”的ONNX模型可以请出性能加速的“王牌”——TensorRT了。TensorRT的工作流程是解析ONNX模型应用一系列优化并为当前GPU硬件生成一个高度定制化的序列化引擎文件.plan或.engine。5.1 使用TensorRT Python API构建引擎我们将编写一个脚本完成从ONNX模型到TensorRT引擎的构建。这个过程可以在部署前离线完成。import tensorrt as trt import os def build_engine_from_onnx(onnx_file_path, engine_file_path, fp16_modeTrue, max_batch_size4): 从ONNX文件构建TensorRT引擎并保存。 参数: onnx_file_path: 输入ONNX模型路径。 engine_file_path: 输出引擎文件路径。 fp16_mode: 是否启用FP16精度能显著提升速度并减少显存占用。 max_batch_size: 最大支持的批次大小需与动态轴设置匹配。 TRT_LOGGER trt.Logger(trt.Logger.WARNING) # 使用WARNING级别减少日志输出 builder trt.Builder(TRT_LOGGER) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, TRT_LOGGER) builder.max_batch_size max_batch_size config builder.create_builder_config() config.max_workspace_size 1 30 # 设置工作空间大小例如1GB if fp16_mode and builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) print(FP16模式已启用。) # 解析ONNX模型 print(f正在从 {onnx_file_path} 解析ONNX模型...) with open(onnx_file_path, rb) as model: if not parser.parse(model.read()): print(ONNX解析失败:) for error in range(parser.num_errors): print(parser.get_error(error)) return None print(ONNX模型解析成功。) # 构建引擎 print(开始构建TensorRT引擎这可能需要几分钟...) engine builder.build_engine(network, config) if engine is None: print(引擎构建失败) return None # 保存引擎到文件 print(f引擎构建成功保存至 {engine_file_path}) with open(engine_file_path, wb) as f: f.write(engine.serialize()) return engine # 使用简化后的ONNX模型构建引擎 onnx_path deoldify_simplified.onnx engine_path deoldify_fp16.engine engine build_engine_from_onnx(onnx_path, engine_path, fp16_modeTrue, max_batch_size4)这个脚本完成了引擎的构建和序列化。生成的后缀为.engine的文件就是优化后的模型可以直接用于高性能推理。5.2 使用构建好的引擎进行推理引擎文件是平台相关的与构建时使用的GPU架构和TensorRT版本绑定。加载引擎进行推理的代码如下import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # 初始化CUDA上下文 import numpy as np class TRTInference: def __init__(self, engine_path): self.TRT_LOGGER trt.Logger(trt.Logger.WARNING) self.engine self.load_engine(engine_path) self.context self.engine.create_execution_context() self.inputs, self.outputs, self.bindings, self.stream self.allocate_buffers() def load_engine(self, engine_path): with open(engine_path, rb) as f, trt.Runtime(self.TRT_LOGGER) as runtime: return runtime.deserialize_cuda_engine(f.read()) def allocate_buffers(self): inputs, outputs, bindings [], [], [] stream cuda.Stream() for binding in self.engine: # 获取绑定信息的形状注意处理动态维度 shape self.engine.get_binding_shape(binding) # 如果维度是动态的-1需要在实际推理前根据输入设置 # 这里假设我们已经知道输入尺寸或者后续会设置 size trt.volume(shape) * self.engine.max_batch_size dtype trt.nptype(self.engine.get_binding_dtype(binding)) # 分配主机和设备内存 host_mem cuda.pagelocked_empty(size, dtype) device_mem cuda.mem_alloc(host_mem.nbytes) bindings.append(int(device_mem)) if self.engine.binding_is_input(binding): inputs.append({host: host_mem, device: device_mem, shape: shape}) else: outputs.append({host: host_mem, device: device_mem, shape: shape}) return inputs, outputs, bindings, stream def infer(self, input_data): # 将输入数据复制到已分配的缓冲区 np.copyto(self.inputs[0][host], input_data.ravel()) # 将输入数据从主机内存传输到设备内存 cuda.memcpy_htod_async(self.inputs[0][device], self.inputs[0][host], self.stream) # 执行推理 self.context.execute_async_v2(bindingsself.bindings, stream_handleself.stream.handle) # 将输出数据从设备内存传输回主机内存 cuda.memcpy_dtoh_async(self.outputs[0][host], self.outputs[0][device], self.stream) # 同步流 self.stream.synchronize() # 返回输出数据需要根据输出形状重塑 output self.outputs[0][host] # 这里需要根据你的模型实际输出形状来重塑例如 [batch, channel, height, width] # 假设输出是单张图像形状为 [1, 3, 256, 256] output_shape (1, 3, 256, 256) # 请替换为你的模型输出形状 return output.reshape(output_shape) # 使用示例 trt_model TRTInference(deoldify_fp16.engine) # 准备一个numpy格式的输入数据例如 shape: (1, 3, 256, 256) dummy_input_np np.random.randn(1, 3, 256, 256).astype(np.float32) output trt_model.infer(dummy_input_np) print(fTensorRT推理完成输出shape: {output.shape})这段代码提供了一个基本的TensorRT推理类。在实际应用中你需要根据模型的输入输出细节特别是动态形状的处理进行调整。6. 性能对比与效果验证转换和加速的最终目的是提升性能。我们用一个简单的测试来对比一下原始PyTorch模型、ONNX Runtime推理以及TensorRT加速后的性能差异。import time import numpy as np def benchmark(model, input_data, warmup10, runs100): 基准测试函数计算平均推理时间 # 预热 for _ in range(warmup): _ model(input_data) # 正式计时 start_time time.time() for _ in range(runs): _ model(input_data) end_time time.time() avg_time (end_time - start_time) * 1000 / runs # 毫秒 return avg_time # 准备测试数据 test_input_np np.random.randn(1, 3, 256, 256).astype(np.float32) # 1. PyTorch 基准测试 (假设model是之前加载的PyTorch模型并在GPU上) test_input_torch torch.from_numpy(test_input_np).cuda() torch_time benchmark(lambda x: model(x), test_input_torch) print(fPyTorch 平均推理时间: {torch_time:.2f} ms) # 2. ONNX Runtime 基准测试 import onnxruntime as ort ort_sess ort.InferenceSession(deoldify_simplified.onnx, providers[CUDAExecutionProvider]) ort_input_name ort_sess.get_inputs()[0].name def ort_infer(x): return ort_sess.run(None, {ort_input_name: x}) ort_time benchmark(ort_infer, test_input_np) print(fONNX Runtime (CUDA) 平均推理时间: {ort_time:.2f} ms) # 3. TensorRT 基准测试 (使用上面定义的TRTInference类) trt_engine TRTInference(deoldify_fp16.engine) def trt_infer(x): return trt_engine.infer(x) trt_time benchmark(trt_infer, test_input_np) print(fTensorRT (FP16) 平均推理时间: {trt_time:.2f} ms) # 计算加速比 print(f\n加速比 (对比PyTorch):) print(f ONNX Runtime: {torch_time/ort_time:.2f}x) print(f TensorRT: {torch_time/trt_time:.2f}x)请注意实际的加速比会因模型结构、输入尺寸、GPU型号、TensorRT配置如是否使用FP16/INT8等因素而有很大差异。对于像DeOldify这样的视觉模型使用TensorRT FP16模式获得2-5倍的加速是比较常见的。如果进一步使用INT8量化速度还可能提升但需要校准数据集并可能带来微小的精度损失。7. 总结与后续建议走完这一整套流程从PyTorch到ONNX再到TensorRT相当于给你的模型做了一次全方位的“性能调优”。最大的感受是前期在模型导出和简化上多花点心思能避免后面很多麻烦。ONNX简化那一步看似简单但经常能解决一些诡异的转换错误。性能提升方面TensorRT确实没让人失望。在我们的测试场景里推理速度的提升是实实在在的。不过也要注意这个优化过程不是全自动的魔法特别是遇到一些特殊算子或者动态形状非常复杂的模型时可能需要手动调整ONNX导出方式或者为TensorRT编写自定义插件Custom Plugin。对于DeOldify或者类似的图像生成、处理模型这条技术路线已经非常成熟了。如果你打算在星图这样的GPU云平台上部署服务强烈建议把生成TensorRT引擎作为部署前的标准步骤。它能显著降低服务延迟提升资源利用率。最后如果遇到问题多看看TensorRT和ONNX的官方文档社区里的讨论也很多。先从简单的、静态输入的模型开始练习熟悉了整个流程和工具链再去挑战更复杂的动态模型会顺利很多。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。