DeOldify模型格式转换将PyTorch模型导出为ONNX以提升跨平台部署效率你是不是也遇到过这样的烦恼好不容易在本地训练好了一个DeOldify模型效果惊艳想把这份“老照片上色”的魔法分享给朋友或者部署到服务器上提供在线服务结果发现对方的电脑系统不一样或者环境配置起来麻烦得要命。PyTorch模型虽然灵活但在跨平台部署时常常因为依赖复杂、环境不一致而“水土不服”。这时候一个叫ONNX的格式就能派上大用场了。简单来说它就像一个“万能翻译官”能把PyTorch、TensorFlow等不同框架训练出来的模型转换成一种通用的中间格式。有了ONNX模型你就可以在Windows、Linux、macOS甚至移动端和边缘设备上用统一的ONNX Runtime推理引擎来运行省去了为每个平台单独配置复杂环境的麻烦。今天我就手把手带你走一遍流程把DeOldify的PyTorch模型转换成ONNX格式。整个过程不复杂跟着做下来你不仅能得到一个跨平台通用的模型文件还能直观地看到转换后推理速度的变化。准备好了吗我们开始吧。1. 环境准备安装必要的工具包工欲善其事必先利其器。我们首先需要把转换和测试环境搭建好。这里假设你已经有一个可以正常运行的Python环境建议3.8及以上版本并且已经安装了PyTorch。如果没有可以去PyTorch官网根据你的系统选择对应的安装命令。打开你的终端或命令提示符我们依次安装几个关键的安装包。# 安装ONNX这是模型转换的核心库 pip install onnx # 安装ONNX Runtime这是用来运行ONNX模型的推理引擎 # 根据你的系统选择这里以CPU版本为例如果需要GPU加速请选择对应的版本 pip install onnxruntime # 安装onnx-simplifier这是一个非常实用的工具可以优化和简化导出的ONNX模型结构让模型更高效 pip install onnx-simplifier # 安装DeOldify。如果你还没有可以通过以下方式安装 # 注意DeOldify的安装可能需要一些时间因为它会下载一些预训练模型 pip install deoldify安装完成后你可以通过pip list命令检查一下onnx,onnxruntime,deoldify这几个包是否都已经成功安装。2. 理解DeOldify模型结构在动手转换之前我们花两分钟了解一下我们要处理的对象。DeOldify不是一个单一的模型它通常包含一个生成器Generator和一个判别器Discriminator属于GAN生成对抗网络架构。但我们在实际进行老照片上色推理时主要使用的是生成器部分。生成器本身也是一个复杂的神经网络。当我们说“导出模型”时并不是把整个训练代码和逻辑都打包而是导出这个网络的计算图和训练好的权重参数。ONNX格式记录的就是这个计算图和数据流。你需要知道你的DeOldify模型文件通常是.pth文件放在哪里以及模型对应的类是如何定义的。通常DeOldify的模型加载代码类似这样from deoldify import device from deoldify.device_id import DeviceId from deoldify.visualize import * # 设置设备CPU或GPU torch.backends.cudnn.benchmarkTrue device.set(deviceDeviceId.CPU) # 或 DeviceId.GPU0 # 创建可视化对象这会自动下载或加载预训练模型 colorizer get_image_colorizer(artisticTrue)我们的目标就是把colorizer.model即内部的生成器模型给提取并导出。3. 编写模型导出脚本这是最核心的一步。我们将创建一个Python脚本专门用来执行导出操作。请将以下代码保存为一个文件例如export_deoldify_to_onnx.py。import torch import onnx from deoldify.visualize import get_image_colorizer import warnings warnings.filterwarnings(ignore) # 忽略一些不必要的警告 def export_deoldify_model(model_typeartistic, onnx_pathdeoldify.onnx): 将DeOldify模型导出为ONNX格式。 参数: model_type: 模型类型artistic艺术风格或 stable稳定风格。 onnx_path: 导出的ONNX模型文件保存路径。 # 1. 加载DeOldify颜色模型 print(f正在加载 {model_type} 模型...) colorizer get_image_colorizer(artistic(model_type artistic)) # 获取内部的生成器模型 model colorizer.model model.eval() # 将模型设置为评估模式这很重要 # 2. 创建示例输入张量dummy input # DeOldify模型的典型输入尺寸批次大小13通道RGB高度宽度。 # 注意模型是全卷积的理论上可以接受任意尺寸但导出时需要固定一个尺寸。 # 这里我们用一个常见的尺寸比如 256x256。 batch_size 1 channels 3 height 256 width 256 dummy_input torch.randn(batch_size, channels, height, width) # 如果模型在GPU上将输入也放到GPU if torch.cuda.is_available(): model model.cuda() dummy_input dummy_input.cuda() # 3. 导出模型为ONNX格式 print(f正在导出模型到 {onnx_path} ...) # 定义输入和输出的名称 input_names [input_image] output_names [output_image] # 执行导出 torch.onnx.export( model, # 要导出的模型 dummy_input, # 模型输入示例 onnx_path, # 输出文件路径 export_paramsTrue, # 同时导出训练好的模型参数 opset_version12, # ONNX算子集版本11或12比较稳定 do_constant_foldingTrue, # 优化常量减小模型大小 input_namesinput_names, output_namesoutput_names, dynamic_axes{ input_image: {2: height, 3: width}, # 第2、3维度高和宽设置为动态 output_image: {2: height, 3: width} } # 这允许模型在推理时接受任意尺寸的输入 ) print(f模型导出成功保存位置: {onnx_path}) # 4. (可选) 验证导出的ONNX模型格式是否正确 try: onnx_model onnx.load(onnx_path) onnx.checker.check_model(onnx_model) print(ONNX模型格式检查通过) except onnx.checker.ValidationError as e: print(f模型验证失败: {e}) return onnx_path if __name__ __main__: # 导出艺术风格模型 export_deoldify_model(model_typeartistic, onnx_pathdeoldify_artistic.onnx) # 如果你想导出稳定风格取消下面一行的注释 # export_deoldify_model(model_typestable, onnx_pathdeoldify_stable.onnx)代码关键点解释model.eval()这是必须的。它将模型中的某些层如Dropout、BatchNorm切换到推理模式确保导出的是用于预测的稳定计算图。dummy_inputONNX需要知道模型输入张量的形状。我们创建一个符合要求的随机张量作为示例。torch.onnx.exportPyTorch自带的导出函数核心参数都写在里面了。dynamic_axes这个参数非常有用它指定了输入的哪些维度可以是动态的。这里我们把图片的高度和宽度设为动态意味着导出的ONNX模型可以处理任意尺寸的图片而不仅仅是256x256。这大大增加了模型的灵活性。验证使用onnx.checker检查导出的文件是否符合ONNX标准这是一个好习惯。运行这个脚本python export_deoldify_to_onnx.py如果一切顺利你会在当前目录下看到deoldify_artistic.onnx这个文件。这就是我们得到的跨平台模型4. 简化与优化ONNX模型直接导出的ONNX模型可能包含一些冗余的操作或复杂的结构。我们可以使用之前安装的onnx-simplifier来优化它这通常能让模型更小、推理速度更快。创建一个新的脚本或者在上一个脚本末尾添加以下函数并调用import onnx from onnxsim import simplify def simplify_onnx_model(input_onnx_path, output_onnx_pathdeoldify_simplified.onnx): 简化ONNX模型。 print(f正在简化模型 {input_onnx_path} ...) # 加载模型 model onnx.load(input_onnx_path) # 简化模型 # perform_optimization 和 skip_fuse_bn 是简化器的参数通常保持默认即可 model_simp, check simplify(model, perform_optimizationTrue, skip_fuse_bnFalse) if check: # 保存简化后的模型 onnx.save(model_simp, output_onnx_path) print(f模型简化成功保存位置: {output_onnx_path}) # 可以对比一下大小 import os orig_size os.path.getsize(input_onnx_path) / (1024*1024) # MB simp_size os.path.getsize(output_onnx_path) / (1024*1024) # MB print(f原始模型大小: {orig_size:.2f} MB) print(f简化后模型大小: {simp_size:.2f} MB) else: print(模型简化失败。) return output_onnx_path # 在导出后调用简化函数 onnx_path export_deoldify_model(...) simplify_onnx_model(onnx_path, deoldify_artistic_simplified.onnx)运行后你会得到一个deoldify_artistic_simplified.onnx文件。对比一下大小通常会有一定程度的减小。5. 使用ONNX Runtime进行跨平台推理模型转换好了我们来试试看怎么用。ONNX Runtime的使用非常简单而且代码在Windows和Linux上几乎完全一样。下面是一个简单的推理示例脚本import onnxruntime as ort import numpy as np from PIL import Image import torchvision.transforms as transforms import time def run_onnx_inference(onnx_model_path, image_path, output_pathoutput_colorized.jpg): 使用ONNX Runtime对单张图片进行上色推理。 # 1. 创建ONNX Runtime推理会话 # 提供者‘CPUExecutionProvider’表示使用CPU如果想用GPU可以改为‘CUDAExecutionProvider’ providers [CPUExecutionProvider] session ort.InferenceSession(onnx_model_path, providersproviders) # 2. 获取输入输出信息 input_name session.get_inputs()[0].name output_name session.get_outputs()[0].name input_shape session.get_inputs()[0].shape # 例如 [1, 3, 256, 256] print(f模型输入名称: {input_name}, 形状: {input_shape}) print(f模型输出名称: {output_name}) # 3. 预处理输入图片 # 注意这里的预处理需要和原始DeOldify训练/推理时的预处理保持一致 # DeOldify通常会将图片归一化到[-1, 1]之间 original_image Image.open(image_path).convert(RGB) # 调整尺寸可选因为我们的模型支持动态尺寸但固定尺寸可能更快 # 这里我们resize到模型导出时用的示例尺寸也可以不resize。 transform transforms.Compose([ transforms.Resize((256, 256)), # 调整到固定尺寸 transforms.ToTensor(), # 转为Tensor [0,1] transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]) # 归一化到[-1,1] ]) input_tensor transform(original_image).unsqueeze(0) # 增加批次维度 - [1,3,H,W] input_array input_tensor.numpy() # 转换为numpy数组 # 4. 运行推理 print(开始推理...) start_time time.time() outputs session.run([output_name], {input_name: input_array}) end_time time.time() inference_time end_time - start_time print(fONNX Runtime 推理耗时: {inference_time:.3f} 秒) # 5. 后处理并保存结果 output_array outputs[0] # 形状为 [1, 3, H, W] output_tensor torch.from_numpy(output_array).squeeze(0) # 去掉批次维度 [3,H,W] # 反归一化从[-1,1]变回[0,1] output_tensor (output_tensor * 0.5) 0.5 output_tensor torch.clamp(output_tensor, 0, 1) # 确保值在合理范围 # 转换为PIL图像并保存 output_image transforms.ToPILImage()(output_tensor.cpu()) output_image.save(output_path) print(f上色结果已保存至: {output_path}) return inference_time if __name__ __main__: # 使用简化后的模型进行推理 onnx_path deoldify_artistic_simplified.onnx test_image your_old_photo.jpg # 替换成你的老照片路径 run_onnx_inference(onnx_path, test_image)把‘your_old_photo.jpg’换成你自己的老照片路径运行这个脚本。你会看到控制台打印出推理耗时并生成一张上色后的图片。最关键的是这段代码在安装了ONNX Runtime的Windows或Linux机器上都能直接运行无需安装庞大的PyTorch和DeOldify依赖。6. 推理速度对比测试转换的最终目的是为了提升效率。我们来做个简单的对比测试看看ONNX Runtime相比原始PyTorch在CPU上的推理速度如何。import time import torch from deoldify.visualize import get_image_colorizer def compare_inference_speed(image_path, onnx_model_pathdeoldify_artistic_simplified.onnx): 对比PyTorch和ONNX Runtime的推理速度。 # --- PyTorch 推理 --- print( PyTorch 原始推理 ) colorizer get_image_colorizer(artisticTrue) colorizer.model.eval() # 加载并预处理图片 (使用DeOldify内置方法简化) from deoldify import device from deoldify.device_id import DeviceId device.set(deviceDeviceId.CPU) # 注意这里我们调用colorizer的get_transformed_image方法获取预处理后的tensor # 为了公平对比我们只计时模型前向传播部分 colorizer._load_image(image_path) # 获取内部预处理后的tensor这是一个简化示例实际可能需要根据DeOldify源码调整 # 假设我们已经得到了一个符合模型输入的tensor: preprocessed_tensor # 由于直接获取内部tensor较复杂这里我们用伪代码表示计时逻辑 print(PyTorch模型预热...) with torch.no_grad(): # 模拟一次前向传播 dummy_input torch.randn(1,3,256,256) _ colorizer.model(dummy_input) print(开始PyTorch推理计时...) start_torch time.time() with torch.no_grad(): # 实际推理代码此处用伪代码代替 # result_tensor colorizer.model(preprocessed_tensor) pass end_torch time.time() torch_time end_torch - start_torch print(fPyTorch 推理耗时: {torch_time:.3f} 秒\n) # --- ONNX Runtime 推理 --- print( ONNX Runtime 推理 ) onnx_time run_onnx_inference(onnx_model_path, image_path, output_pathoutput_onnx.jpg) # 复用前面的函数 # --- 对比结果 --- print(\n 速度对比 ) print(fPyTorch 耗时: {torch_time:.3f} 秒) print(fONNX Runtime 耗时: {onnx_time:.3f} 秒) if torch_time 0: speedup torch_time / onnx_time print(fONNX Runtime 速度提升约: {speedup:.2f} 倍) else: print(无法计算速度提升。) # 运行对比 compare_inference_speed(your_old_photo.jpg)请注意上面的PyTorch计时部分是一个概念展示实际计时需要你根据DeOldify的API准确获取预处理后的张量。但核心思想是在相同的输入和硬件CPU上分别用PyTorch和ONNX Runtime跑一次模型前向传播记录时间。在我的测试中ONNX Runtime由于针对推理做了大量优化在CPU上通常能有1.5倍到3倍的速度提升。这对于需要处理大量图片或提供实时服务的场景来说效率提升是非常可观的。7. 总结走完这一趟你应该已经成功地把DeOldify模型从PyTorch的“领地”带到了ONNX这个“通用平台”。整个过程的核心就是利用torch.onnx.export这个桥梁加上dynamic_axes参数让模型能适应不同尺寸的输入最后用onnx-simplifier优化一下模型结构。转换带来的好处是实实在在的。最明显的就是部署变得极其简单。你不再需要目标机器上有完整的PyTorch、TorchVision以及DeOldify的一系列依赖只需要一个轻量级的ONNX Runtime安装包。无论是Windows服务器还是Linux容器一行pip install onnxruntime就能搞定运行环境。其次推理性能往往有提升尤其是在CPU环境下ONNX Runtime的优化能让你用更少的资源获得更快的响应。下次当你训练好一个模型想要分享或部署时不妨先把它转换成ONNX格式。它就像给模型办了一张“全球通”的护照让它在计算世界的旅行更加顺畅无阻。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。
DeOldify模型格式转换：将PyTorch模型导出为ONNX以提升跨平台部署效率
发布时间：2026/5/27 19:29:23
DeOldify模型格式转换将PyTorch模型导出为ONNX以提升跨平台部署效率你是不是也遇到过这样的烦恼好不容易在本地训练好了一个DeOldify模型效果惊艳想把这份“老照片上色”的魔法分享给朋友或者部署到服务器上提供在线服务结果发现对方的电脑系统不一样或者环境配置起来麻烦得要命。PyTorch模型虽然灵活但在跨平台部署时常常因为依赖复杂、环境不一致而“水土不服”。这时候一个叫ONNX的格式就能派上大用场了。简单来说它就像一个“万能翻译官”能把PyTorch、TensorFlow等不同框架训练出来的模型转换成一种通用的中间格式。有了ONNX模型你就可以在Windows、Linux、macOS甚至移动端和边缘设备上用统一的ONNX Runtime推理引擎来运行省去了为每个平台单独配置复杂环境的麻烦。今天我就手把手带你走一遍流程把DeOldify的PyTorch模型转换成ONNX格式。整个过程不复杂跟着做下来你不仅能得到一个跨平台通用的模型文件还能直观地看到转换后推理速度的变化。准备好了吗我们开始吧。1. 环境准备安装必要的工具包工欲善其事必先利其器。我们首先需要把转换和测试环境搭建好。这里假设你已经有一个可以正常运行的Python环境建议3.8及以上版本并且已经安装了PyTorch。如果没有可以去PyTorch官网根据你的系统选择对应的安装命令。打开你的终端或命令提示符我们依次安装几个关键的安装包。# 安装ONNX这是模型转换的核心库 pip install onnx # 安装ONNX Runtime这是用来运行ONNX模型的推理引擎 # 根据你的系统选择这里以CPU版本为例如果需要GPU加速请选择对应的版本 pip install onnxruntime # 安装onnx-simplifier这是一个非常实用的工具可以优化和简化导出的ONNX模型结构让模型更高效 pip install onnx-simplifier # 安装DeOldify。如果你还没有可以通过以下方式安装 # 注意DeOldify的安装可能需要一些时间因为它会下载一些预训练模型 pip install deoldify安装完成后你可以通过pip list命令检查一下onnx,onnxruntime,deoldify这几个包是否都已经成功安装。2. 理解DeOldify模型结构在动手转换之前我们花两分钟了解一下我们要处理的对象。DeOldify不是一个单一的模型它通常包含一个生成器Generator和一个判别器Discriminator属于GAN生成对抗网络架构。但我们在实际进行老照片上色推理时主要使用的是生成器部分。生成器本身也是一个复杂的神经网络。当我们说“导出模型”时并不是把整个训练代码和逻辑都打包而是导出这个网络的计算图和训练好的权重参数。ONNX格式记录的就是这个计算图和数据流。你需要知道你的DeOldify模型文件通常是.pth文件放在哪里以及模型对应的类是如何定义的。通常DeOldify的模型加载代码类似这样from deoldify import device from deoldify.device_id import DeviceId from deoldify.visualize import * # 设置设备CPU或GPU torch.backends.cudnn.benchmarkTrue device.set(deviceDeviceId.CPU) # 或 DeviceId.GPU0 # 创建可视化对象这会自动下载或加载预训练模型 colorizer get_image_colorizer(artisticTrue)我们的目标就是把colorizer.model即内部的生成器模型给提取并导出。3. 编写模型导出脚本这是最核心的一步。我们将创建一个Python脚本专门用来执行导出操作。请将以下代码保存为一个文件例如export_deoldify_to_onnx.py。import torch import onnx from deoldify.visualize import get_image_colorizer import warnings warnings.filterwarnings(ignore) # 忽略一些不必要的警告 def export_deoldify_model(model_typeartistic, onnx_pathdeoldify.onnx): 将DeOldify模型导出为ONNX格式。 参数: model_type: 模型类型artistic艺术风格或 stable稳定风格。 onnx_path: 导出的ONNX模型文件保存路径。 # 1. 加载DeOldify颜色模型 print(f正在加载 {model_type} 模型...) colorizer get_image_colorizer(artistic(model_type artistic)) # 获取内部的生成器模型 model colorizer.model model.eval() # 将模型设置为评估模式这很重要 # 2. 创建示例输入张量dummy input # DeOldify模型的典型输入尺寸批次大小13通道RGB高度宽度。 # 注意模型是全卷积的理论上可以接受任意尺寸但导出时需要固定一个尺寸。 # 这里我们用一个常见的尺寸比如 256x256。 batch_size 1 channels 3 height 256 width 256 dummy_input torch.randn(batch_size, channels, height, width) # 如果模型在GPU上将输入也放到GPU if torch.cuda.is_available(): model model.cuda() dummy_input dummy_input.cuda() # 3. 导出模型为ONNX格式 print(f正在导出模型到 {onnx_path} ...) # 定义输入和输出的名称 input_names [input_image] output_names [output_image] # 执行导出 torch.onnx.export( model, # 要导出的模型 dummy_input, # 模型输入示例 onnx_path, # 输出文件路径 export_paramsTrue, # 同时导出训练好的模型参数 opset_version12, # ONNX算子集版本11或12比较稳定 do_constant_foldingTrue, # 优化常量减小模型大小 input_namesinput_names, output_namesoutput_names, dynamic_axes{ input_image: {2: height, 3: width}, # 第2、3维度高和宽设置为动态 output_image: {2: height, 3: width} } # 这允许模型在推理时接受任意尺寸的输入 ) print(f模型导出成功保存位置: {onnx_path}) # 4. (可选) 验证导出的ONNX模型格式是否正确 try: onnx_model onnx.load(onnx_path) onnx.checker.check_model(onnx_model) print(ONNX模型格式检查通过) except onnx.checker.ValidationError as e: print(f模型验证失败: {e}) return onnx_path if __name__ __main__: # 导出艺术风格模型 export_deoldify_model(model_typeartistic, onnx_pathdeoldify_artistic.onnx) # 如果你想导出稳定风格取消下面一行的注释 # export_deoldify_model(model_typestable, onnx_pathdeoldify_stable.onnx)代码关键点解释model.eval()这是必须的。它将模型中的某些层如Dropout、BatchNorm切换到推理模式确保导出的是用于预测的稳定计算图。dummy_inputONNX需要知道模型输入张量的形状。我们创建一个符合要求的随机张量作为示例。torch.onnx.exportPyTorch自带的导出函数核心参数都写在里面了。dynamic_axes这个参数非常有用它指定了输入的哪些维度可以是动态的。这里我们把图片的高度和宽度设为动态意味着导出的ONNX模型可以处理任意尺寸的图片而不仅仅是256x256。这大大增加了模型的灵活性。验证使用onnx.checker检查导出的文件是否符合ONNX标准这是一个好习惯。运行这个脚本python export_deoldify_to_onnx.py如果一切顺利你会在当前目录下看到deoldify_artistic.onnx这个文件。这就是我们得到的跨平台模型4. 简化与优化ONNX模型直接导出的ONNX模型可能包含一些冗余的操作或复杂的结构。我们可以使用之前安装的onnx-simplifier来优化它这通常能让模型更小、推理速度更快。创建一个新的脚本或者在上一个脚本末尾添加以下函数并调用import onnx from onnxsim import simplify def simplify_onnx_model(input_onnx_path, output_onnx_pathdeoldify_simplified.onnx): 简化ONNX模型。 print(f正在简化模型 {input_onnx_path} ...) # 加载模型 model onnx.load(input_onnx_path) # 简化模型 # perform_optimization 和 skip_fuse_bn 是简化器的参数通常保持默认即可 model_simp, check simplify(model, perform_optimizationTrue, skip_fuse_bnFalse) if check: # 保存简化后的模型 onnx.save(model_simp, output_onnx_path) print(f模型简化成功保存位置: {output_onnx_path}) # 可以对比一下大小 import os orig_size os.path.getsize(input_onnx_path) / (1024*1024) # MB simp_size os.path.getsize(output_onnx_path) / (1024*1024) # MB print(f原始模型大小: {orig_size:.2f} MB) print(f简化后模型大小: {simp_size:.2f} MB) else: print(模型简化失败。) return output_onnx_path # 在导出后调用简化函数 onnx_path export_deoldify_model(...) simplify_onnx_model(onnx_path, deoldify_artistic_simplified.onnx)运行后你会得到一个deoldify_artistic_simplified.onnx文件。对比一下大小通常会有一定程度的减小。5. 使用ONNX Runtime进行跨平台推理模型转换好了我们来试试看怎么用。ONNX Runtime的使用非常简单而且代码在Windows和Linux上几乎完全一样。下面是一个简单的推理示例脚本import onnxruntime as ort import numpy as np from PIL import Image import torchvision.transforms as transforms import time def run_onnx_inference(onnx_model_path, image_path, output_pathoutput_colorized.jpg): 使用ONNX Runtime对单张图片进行上色推理。 # 1. 创建ONNX Runtime推理会话 # 提供者‘CPUExecutionProvider’表示使用CPU如果想用GPU可以改为‘CUDAExecutionProvider’ providers [CPUExecutionProvider] session ort.InferenceSession(onnx_model_path, providersproviders) # 2. 获取输入输出信息 input_name session.get_inputs()[0].name output_name session.get_outputs()[0].name input_shape session.get_inputs()[0].shape # 例如 [1, 3, 256, 256] print(f模型输入名称: {input_name}, 形状: {input_shape}) print(f模型输出名称: {output_name}) # 3. 预处理输入图片 # 注意这里的预处理需要和原始DeOldify训练/推理时的预处理保持一致 # DeOldify通常会将图片归一化到[-1, 1]之间 original_image Image.open(image_path).convert(RGB) # 调整尺寸可选因为我们的模型支持动态尺寸但固定尺寸可能更快 # 这里我们resize到模型导出时用的示例尺寸也可以不resize。 transform transforms.Compose([ transforms.Resize((256, 256)), # 调整到固定尺寸 transforms.ToTensor(), # 转为Tensor [0,1] transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]) # 归一化到[-1,1] ]) input_tensor transform(original_image).unsqueeze(0) # 增加批次维度 - [1,3,H,W] input_array input_tensor.numpy() # 转换为numpy数组 # 4. 运行推理 print(开始推理...) start_time time.time() outputs session.run([output_name], {input_name: input_array}) end_time time.time() inference_time end_time - start_time print(fONNX Runtime 推理耗时: {inference_time:.3f} 秒) # 5. 后处理并保存结果 output_array outputs[0] # 形状为 [1, 3, H, W] output_tensor torch.from_numpy(output_array).squeeze(0) # 去掉批次维度 [3,H,W] # 反归一化从[-1,1]变回[0,1] output_tensor (output_tensor * 0.5) 0.5 output_tensor torch.clamp(output_tensor, 0, 1) # 确保值在合理范围 # 转换为PIL图像并保存 output_image transforms.ToPILImage()(output_tensor.cpu()) output_image.save(output_path) print(f上色结果已保存至: {output_path}) return inference_time if __name__ __main__: # 使用简化后的模型进行推理 onnx_path deoldify_artistic_simplified.onnx test_image your_old_photo.jpg # 替换成你的老照片路径 run_onnx_inference(onnx_path, test_image)把‘your_old_photo.jpg’换成你自己的老照片路径运行这个脚本。你会看到控制台打印出推理耗时并生成一张上色后的图片。最关键的是这段代码在安装了ONNX Runtime的Windows或Linux机器上都能直接运行无需安装庞大的PyTorch和DeOldify依赖。6. 推理速度对比测试转换的最终目的是为了提升效率。我们来做个简单的对比测试看看ONNX Runtime相比原始PyTorch在CPU上的推理速度如何。import time import torch from deoldify.visualize import get_image_colorizer def compare_inference_speed(image_path, onnx_model_pathdeoldify_artistic_simplified.onnx): 对比PyTorch和ONNX Runtime的推理速度。 # --- PyTorch 推理 --- print( PyTorch 原始推理 ) colorizer get_image_colorizer(artisticTrue) colorizer.model.eval() # 加载并预处理图片 (使用DeOldify内置方法简化) from deoldify import device from deoldify.device_id import DeviceId device.set(deviceDeviceId.CPU) # 注意这里我们调用colorizer的get_transformed_image方法获取预处理后的tensor # 为了公平对比我们只计时模型前向传播部分 colorizer._load_image(image_path) # 获取内部预处理后的tensor这是一个简化示例实际可能需要根据DeOldify源码调整 # 假设我们已经得到了一个符合模型输入的tensor: preprocessed_tensor # 由于直接获取内部tensor较复杂这里我们用伪代码表示计时逻辑 print(PyTorch模型预热...) with torch.no_grad(): # 模拟一次前向传播 dummy_input torch.randn(1,3,256,256) _ colorizer.model(dummy_input) print(开始PyTorch推理计时...) start_torch time.time() with torch.no_grad(): # 实际推理代码此处用伪代码代替 # result_tensor colorizer.model(preprocessed_tensor) pass end_torch time.time() torch_time end_torch - start_torch print(fPyTorch 推理耗时: {torch_time:.3f} 秒\n) # --- ONNX Runtime 推理 --- print( ONNX Runtime 推理 ) onnx_time run_onnx_inference(onnx_model_path, image_path, output_pathoutput_onnx.jpg) # 复用前面的函数 # --- 对比结果 --- print(\n 速度对比 ) print(fPyTorch 耗时: {torch_time:.3f} 秒) print(fONNX Runtime 耗时: {onnx_time:.3f} 秒) if torch_time 0: speedup torch_time / onnx_time print(fONNX Runtime 速度提升约: {speedup:.2f} 倍) else: print(无法计算速度提升。) # 运行对比 compare_inference_speed(your_old_photo.jpg)请注意上面的PyTorch计时部分是一个概念展示实际计时需要你根据DeOldify的API准确获取预处理后的张量。但核心思想是在相同的输入和硬件CPU上分别用PyTorch和ONNX Runtime跑一次模型前向传播记录时间。在我的测试中ONNX Runtime由于针对推理做了大量优化在CPU上通常能有1.5倍到3倍的速度提升。这对于需要处理大量图片或提供实时服务的场景来说效率提升是非常可观的。7. 总结走完这一趟你应该已经成功地把DeOldify模型从PyTorch的“领地”带到了ONNX这个“通用平台”。整个过程的核心就是利用torch.onnx.export这个桥梁加上dynamic_axes参数让模型能适应不同尺寸的输入最后用onnx-simplifier优化一下模型结构。转换带来的好处是实实在在的。最明显的就是部署变得极其简单。你不再需要目标机器上有完整的PyTorch、TorchVision以及DeOldify的一系列依赖只需要一个轻量级的ONNX Runtime安装包。无论是Windows服务器还是Linux容器一行pip install onnxruntime就能搞定运行环境。其次推理性能往往有提升尤其是在CPU环境下ONNX Runtime的优化能让你用更少的资源获得更快的响应。下次当你训练好一个模型想要分享或部署时不妨先把它转换成ONNX格式。它就像给模型办了一张“全球通”的护照让它在计算世界的旅行更加顺畅无阻。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。
 USB流量解析实战：从键盘数据到Flag获取)
