cv_unet_image-colorization模型转换与部署从PyTorch到ONNX最近有不少朋友在尝试用深度学习给黑白照片上色效果确实挺有意思的。但大家普遍遇到一个问题训练好的PyTorch模型怎么才能方便地用在各种不同的地方比如一个简单的Web应用里或者直接集成到手机App里总不能每次都带着一整套PyTorch环境跑吧。这就涉及到模型转换和部署了。今天我就以经典的图像上色模型cv_unet_image-colorization为例带你走一遍从PyTorch模型到ONNX格式再到用ONNX Runtime进行推理的完整流程。整个过程其实不难关键是把几个步骤理清楚。我会尽量用大白话讲明白并提供可以直接运行的代码让你看完就能动手实践。1. 准备工作与环境搭建在开始转换之前我们得先把“厨房”收拾好把需要的“食材”和“工具”备齐。这里没什么高深的理论就是一些必要的软件包。1.1 核心工具包安装你需要一个Python环境建议3.8或以上版本然后通过pip安装以下几个核心库。打开你的终端或命令行一行命令搞定pip install torch torchvision onnx onnxruntime pillow numpy opencv-python简单解释一下这几个包是干嘛的torchtorchvision: 这是PyTorch的本体用来加载和验证我们训练好的原始模型。onnx: 这是ONNX格式的Python接口负责执行模型转换。onnxruntime: 这是微软推出的高性能推理引擎我们转换后的模型就靠它来运行。pillowopencv-python: 这两个是图像处理的常用库用来读图片、调整大小、转换颜色空间等。numpy: 基础的科学计算库处理数组数据。如果你的电脑有NVIDIA显卡并且想用GPU来加速推理可以额外安装onnxruntime-gpupip install onnxruntime-gpu注意安装onnxruntime-gpu需要你的CUDA版本和它匹配。如果遇到问题先只用CPU版本的onnxruntime也是完全可行的我们后面会对比两者的速度。1.2 获取模型与理解输入输出工欲善其事必先利其器。我们得先拿到要转换的模型。cv_unet_image-colorization模型通常是一个PyTorch的.pth文件。你可以从原始的训练代码仓库下载或者如果你自己训练过就用你自己的模型文件。关键一步弄清楚模型的“胃口”和“产出”。在转换前你必须明确知道你的模型输入Input它吃进去的图片是什么样子的比如尺寸是256x256还是512x512是灰度图单通道还是RGB图三通道数值范围是0-255还是0-1输出Output它吐出来的是什么对于上色模型通常是上色后的图片比如Lab色彩空间下的ab通道。如果不清楚最好的办法是回头看看模型的训练代码或者加载模型后打印一下结构。假设我们的模型colorization_model.pth需要输入一张256x256的灰度图L通道并输出对应的ab颜色通道。2. 核心步骤PyTorch模型转ONNX准备好了环境和模型现在进入正题——转换。这个过程就像把一份中文菜谱PyTorch翻译成一种世界通用的食谱格式ONNX这样任何能看懂这种通用格式的“厨师”推理引擎都能照着做。2.1 加载PyTorch模型首先我们用PyTorch把原来的模型加载进来。这里假设模型定义在一个叫UNet的类里。import torch import torch.nn as nn # 假设你的模型定义类如下需要与你训练时的定义一致 class UNet(nn.Module): def __init__(self): super(UNet, self).__init__() # ... 这里是你模型的实际结构定义 # 例如编码器、解码器、跳跃连接等 self.encoder ... self.decoder ... def forward(self, x): # ... 前向传播逻辑 return x # 实例化模型并加载权重 model UNet() model.load_state_dict(torch.load(colorization_model.pth, map_locationcpu)) # 将模型设置为评估模式这很重要会关闭Dropout、BatchNorm的统计更新等训练特有的行为。 model.eval()重要提示model.eval()这一行必不可少。它告诉模型“现在是考试不是练习”会固定住一些层如BatchNorm的运行状态确保转换和推理时的一致性。2.2 创建虚拟输入并执行转换ONNX转换需要一个“例子输入”用来追踪模型的计算图。我们根据之前了解的输入要求创建一个符合形状和类型的虚拟张量。import torch # 创建一个虚拟输入张量 # 假设输入是 [batch_size, channels, height, width] # 对于灰度图channels1对于RGBchannels3 # 这里我们假设是单张256x256的灰度图 batch_size 1 channels 1 height 256 width 256 dummy_input torch.randn(batch_size, channels, height, width) # 指定导出的ONNX文件名 onnx_model_path colorization_model.onnx # 执行转换 torch.onnx.export( model, # 要转换的模型 dummy_input, # 模型输入例子 onnx_model_path, # 输出ONNX文件路径 export_paramsTrue, # 将模型参数权重也保存在文件中 opset_version12, # ONNX算子集版本11或12比较通用稳定 do_constant_foldingTrue, # 优化常量减小模型大小 input_names[input], # 输入节点的名称 output_names[output], # 输出节点的名称 dynamic_axes{ # 指定动态维度可选用于支持可变输入尺寸 input: {0: batch_size, 2: height, 3: width}, output: {0: batch_size, 2: height, 3: width} } ) print(f模型已成功导出至: {onnx_model_path})参数解释opset_version: 建议设为11或12兼容性较好。dynamic_axes: 这个参数非常有用。它允许你的模型在推理时接受不同尺寸的图片比如不是固定的256x256。上面代码中{0: batch_size, 2: height, 3: width}表示第0维批量、第2维高、第3维宽是动态的。如果你只需要固定尺寸可以省略这个参数。2.3 验证转换结果转换完成后别急着用先验货。我们可以用ONNX的Python包简单检查一下模型格式是否正确并看看它的输入输出信息。import onnx # 加载并检查ONNX模型 onnx_model onnx.load(onnx_model_path) onnx.checker.check_model(onnx_model) print(ONNX模型格式检查通过) # 打印模型输入输出信息 print( 模型摘要 ) for input in onnx_model.graph.input: print(f输入名称: {input.name}, 形状: {input.type.tensor_type.shape}) for output in onnx_model.graph.output: print(f输出名称: {output.name}, 形状: {output.type.tensor_type.shape})如果看到打印出的输入输出形状和你预期的一致比如输入是[1, 1, 256, 256]那说明转换基本成功了。3. 使用ONNX Runtime进行推理模型转换好了现在我们来“尝尝”用它做推理的滋味。ONNX Runtime的使用非常直观。3.1 基础推理流程我们先写一个完整的函数实现从读取图片、预处理、推理到后处理保存的流程。import onnxruntime as ort import numpy as np from PIL import Image import cv2 def colorize_image_onnx(image_path, onnx_model_path, output_path): 使用ONNX模型为灰度图上色 Args: image_path: 输入灰度图路径 onnx_model_path: ONNX模型文件路径 output_path: 输出彩色图路径 # 1. 创建ONNX Runtime推理会话 # 默认使用CPU。如果想用GPUproviders参数可以改为[CUDAExecutionProvider] session ort.InferenceSession(onnx_model_path, providers[CPUExecutionProvider]) # 2. 读取并预处理图像 # 假设模型需要256x256的灰度图单通道值范围0-1 img Image.open(image_path).convert(L) # 转换为灰度图 img img.resize((256, 256)) # 调整尺寸 img_array np.array(img, dtypenp.float32) / 255.0 # 归一化到[0,1] # 调整维度为 [1, 1, H, W] (NCHW格式) input_tensor img_array[np.newaxis, np.newaxis, :, :] # 3. 运行推理 # session.run([输出名称列表], {输入名称: 输入数据}) outputs session.run([output], {input: input_tensor}) # outputs是一个列表我们取第一个输出我们的ab通道预测 ab_channels outputs[0] # 假设形状是[1, 2, 256, 256] # 4. 后处理将预测的ab通道与原始L通道合并并转换回RGB # 注意这里需要根据你模型的具体输出进行适配。 # 假设ab_channels范围是[-1, 1]需要先调整。 ab_channels (ab_channels.squeeze() * 127).astype(np.uint8) # 缩放到0-255范围并去除批次维度 L_channel (img_array * 255).astype(np.uint8) # 合并Lab通道 (OpenCV使用BGR顺序但Lab是相同的) lab_image cv2.merge([L_channel, ab_channels[0], ab_channels[1]]) # 将Lab转换回BGR rgb_image cv2.cvtColor(lab_image, cv2.COLOR_LAB2BGR) # 再转换到RGB以供PIL保存 rgb_image cv2.cvtColor(rgb_image, cv2.COLOR_BGR2RGB) # 5. 保存结果 result_img Image.fromarray(rgb_image) result_img.save(output_path) print(f上色完成结果已保存至: {output_path}) return rgb_image # 使用示例 if __name__ __main__: colorize_image_onnx(old_photo.jpg, colorization_model.onnx, colorized_photo.jpg)关键点ort.InferenceSession: 这是推理的核心对象。创建它时会加载并优化模型。session.run: 执行推理。你需要提供输出节点的名称列表和输入数据的字典。后处理这部分代码需要根据你的模型具体输出进行调整。上面的例子是一个典型的Lab色彩空间上色流程。如果你的模型输出直接是RGB或者数值范围不同处理方式会不一样。务必与你训练/验证时的后处理代码保持一致。3.2 性能对比CPU vs GPU大家最关心的问题之一就是用GPU到底能快多少我们来简单测一下。确保你安装了onnxruntime-gpu并且CUDA可用。import time import onnxruntime as ort def benchmark_inference(onnx_model_path, use_gpuFalse, warmup5, runs20): 基准测试推理速度 providers [CUDAExecutionProvider, CPUExecutionProvider] if use_gpu else [CPUExecutionProvider] session ort.InferenceSession(onnx_model_path, providersproviders) # 准备一个随机输入数据 dummy_input np.random.randn(1, 1, 256, 256).astype(np.float32) # 预热 for _ in range(warmup): _ session.run([output], {input: dummy_input}) # 正式计时 start_time time.time() for _ in range(runs): _ session.run([output], {input: dummy_input}) end_time time.time() avg_time (end_time - start_time) / runs * 1000 # 平均每次推理的毫秒数 fps 1000 / avg_time if avg_time 0 else 0 # 每秒帧数 device GPU if use_gpu else CPU print(f[{device}] 平均推理时间: {avg_time:.2f} ms, 预估FPS: {fps:.2f}) return avg_time, fps # 分别测试CPU和GPU print(开始性能基准测试...) cpu_time, cpu_fps benchmark_inference(colorization_model.onnx, use_gpuFalse) gpu_time, gpu_fps benchmark_inference(colorization_model.onnx, use_gpuTrue) print(f\n性能对比) print(fGPU速度是CPU的 {cpu_time/gpu_time:.2f} 倍)跑一下这个脚本你就能直观地看到在你自己机器上的加速效果。一般来说对于UNet这类卷积网络GPU能有几倍到几十倍的提升尤其是处理批量图片时优势更明显。4. 跨平台部署与实用技巧把模型转换成ONNX最大的好处就是“一次转换到处运行”。下面看看几种常见的部署场景。4.1 在Web后端服务中部署你可以用Flask、FastAPI等框架轻松创建一个API服务。# 这是一个使用FastAPI的简单示例 from fastapi import FastAPI, File, UploadFile from fastapi.responses import FileResponse import numpy as np import onnxruntime as ort from PIL import Image import io import cv2 app FastAPI() # 在服务启动时加载模型避免每次请求都加载 session ort.InferenceSession(colorization_model.onnx, providers[CPUExecutionProvider]) app.post(/colorize/) async def colorize_image(file: UploadFile File(...)): # 1. 读取上传的图片 contents await file.read() image Image.open(io.BytesIO(contents)).convert(L) image image.resize((256, 256)) input_array np.array(image, dtypenp.float32) / 255.0 input_tensor input_array[np.newaxis, np.newaxis, :, :] # 2. 推理 outputs session.run([output], {input: input_tensor}) ab_channels outputs[0] # ... (后处理代码同上例) # 3. 将结果图像保存到字节流 result_image_pil Image.fromarray(rgb_image) img_byte_arr io.BytesIO() result_image_pil.save(img_byte_arr, formatJPEG) img_byte_arr.seek(0) # 4. 返回图像 return FileResponse(img_byte_arr, media_typeimage/jpeg, filenamecolorized.jpg) # 运行: uvicorn main:app --reload这样前端页面就可以通过上传图片到这个API端点来获取上色后的结果了。4.2 模型优化与注意事项为了让部署更顺畅这里还有几个小贴士模型简化与量化ONNX Runtime支持模型量化将浮点数权重转换为8位整数能显著减小模型体积并提升推理速度尤其适合移动端和边缘设备。可以使用onnxruntime的量化工具进行操作。输入尺寸固定化虽然我们之前设置了动态尺寸但固定输入尺寸通常能获得更好的性能优化。如果你确定输入尺寸不变可以在导出ONNX时不设置dynamic_axes或者使用ONNX Runtime的图优化功能固定尺寸。错误处理在生产环境中务必在代码中添加健壮的错误处理比如检查输入图片格式、模型加载失败、推理出错等情况。版本一致性确保训练、转换、推理各环节使用的库版本尤其是PyTorch和ONNX没有大的兼容性问题。有时候版本不匹配会导致奇怪的错误。5. 总结走完这一趟你会发现把PyTorch模型转换成ONNX并部署起来并没有想象中那么复杂。核心就是三步用torch.onnx.export转换模型用onnxruntime.InferenceSession加载运行再配上合适的前后处理。最大的收获是灵活性。ONNX模型就像一份通用说明书可以被很多不同的“机器”推理引擎读懂比如ONNX Runtime、TensorRT、OpenVINO等等。这意味着你今天在Python服务里用的模型明天可以几乎不改动地放到C程序里或者通过其他支持ONNX的框架在移动端运行。实际用下来转换过程最需要注意的就是输入输出格式的对应以及后处理逻辑要和训练时保持一致。性能方面如果条件允许用GPU推理带来的速度提升是实实在在的对于需要实时处理或者批量处理的场景非常有用。如果你刚开始接触模型部署可以从这个小项目练手。成功跑通后再尝试更复杂的模型和优化技巧比如量化、多线程推理等路就会越走越宽了。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。