PyTorch模型转ONNX及部署优化指南

将PyTorch模型转换为ONNX以便通过FastAPI（Docker）部署

在本课中，将学习如何使用PyTorch Image Models（TIMM）将预训练的ResNetV2-50模型转换为ONNX，分析其结构，并使用ONNX Runtime测试推理。同时，将对比ONNX与标准PyTorch执行在推理速度和模型大小上的差异，以阐明为何ONNX更适合轻量级AI推理。这为模型与FastAPI和Docker的集成做好了准备，确保了在部署到某中心Lambda之前的环境一致性。

本课是某中心Lambda系列4部分课程中的第2课。

引言

在生产环境中部署AI模型需要快速、可扩展且高效的推理。然而，当迁移到某中心Lambda等云环境时，需要考虑几个限制因素：

  • 有限内存：某中心Lambda最多支持10GB内存，这对于大型模型可能不够。
  • 执行时间限制：单个Lambda函数执行时间不能超过15分钟。
  • 冷启动：如果Lambda函数闲置一段时间，启动时间会变长，从而降低推理速度。
  • 存储限制：Lambda包有250MB的限制（包括依赖项），因此优化模型大小至关重要。

为确保顺利部署和高效推理，我们必须在将AI模型集成到FastAPI之前，先在本地进行测试和优化。这意味着：

  1. 将模型转换为轻量级、优化的格式（ONNX）。
  2. 比较ONNX与PyTorch，分析在速度、性能和文件大小上的差异。
  3. 通过API（FastAPI）提供服务，以便轻松处理推理请求。
  4. 使用Docker对设置进行容器化，以确保跨环境的一致执行。

回顾上一课

在第1课中，我们探讨了使用某中心Lambda、API Gateway和ONNX Runtime进行无服务器AI推理的基础知识，为什么ONNX是资源受限环境中AI推理的理想选择，某中心Lambda如何处理无服务器设置中的推理请求，以及如何设置开发环境以使用ONNX和FastAPI。

现在，在本课中，我们将通过以下步骤继续：

  • 将PyTorch模型转换为ONNX。
  • 在推理速度和模型大小方面比较ONNX和PyTorch。
  • 使用FastAPI服务模型并在Docker中运行。

为什么这一步很重要

许多开发者在没有在受控环境中进行适当测试的情况下，就试图直接将AI模型部署到某中心Lambda，这通常会导致：

  • 由于兼容性问题导致模型执行错误。
  • 在云环境中推理延迟高。
  • 内存使用效率低下导致功能失败。

我们首先使用ONNX Runtime和FastAPI在本地构建、测试和优化推理服务器，而不是过早部署。这使我们能够：

  • 通过减少大小和提高推理速度来优化模型以进行部署。
  • 比较ONNX与PyTorch，以确定某中心Lambda的最佳格式。
  • 在本地捕获和调试错误，而不是处理某中心上的部署故障。
  • 使用Docker保持环境一致，确保本地有效的内容在生产中也有效。

本课学习内容

在本课中，我们将从以下内容开始：

  • 将PyTorch ResNetV2-50（TIMM）模型转换为ONNX格式以进行高效推理。
  • 分析ONNX模型结构，了解其与PyTorch的区别。
  • 比较ONNX与PyTorch的推理性能，包括执行时间和模型大小。
  • 使用ONNX Runtime验证推理结果以确保准确性。

完成后，在下一课中，我们将：

  • 构建基于FastAPI的推理API来处理图像分类请求。
  • 在部署前本地测试API。
  • 使用Docker对FastAPI服务器进行容器化以实现一致的运行时。

使用的技术

  • ONNX Runtime：ONNX是一个用于机器学习模型的开放标准，它允许我们：优化模型在不同硬件架构上的执行；与原生PyTorch或TensorFlow模型相比，减少推理延迟；使模型在不同云和边缘环境中更具可移植性。
  • PyTorch Image Models（TIMM）：我们将利用TIMM（Torch Image Models）——一个提供预训练模型和高效实现的强大库。TIMM使得：加载和微调最先进的图像分类模型；以最小的努力将模型转换为ONNX格式；确保与ONNX Runtime的兼容性以进行优化推理变得更加容易。
  • FastAPI（下一课涵盖）：FastAPI是一个高性能Web框架，它允许我们：将ONNX模型作为推理API公开；以最小的开销高效处理实时请求；通过保持API轻量级来确保后续顺利过渡到某中心Lambda。
  • Docker（下一课涵盖）：对FastAPI推理服务器进行容器化可确保：在不同环境（本地、某中心Lambda、云、边缘）中的一致执行；依赖项隔离可防止系统包之间的冲突；简化部署，因为某中心Lambda支持容器化应用程序。

配置开发环境

为了确保开发过程顺利进行，我们首先需要使用所需的依赖项来设置环境。这包括安装FastAPI、ONNX Runtime、PyTorch Image Models（TIMM）和Docker，以便在部署到某中心Lambda之前，在容器内运行推理API。

运行以下命令安装所需的Python包：

$ pip install fastapi[all]==0.98.0 numpy==1.25.2 onnxruntime==1.15.1 mangum==0.17.0 Pillow==9.5.0 timm==0.9.5 onnx==1.14.0

项目结构

首先需要查看项目目录结构。

$ tree . -L 1
resnet-aws-serverless-classifier/
├── src/
│   ├── convert.py           # 此处使用 – 将PyTorch模型转换为ONNX
│   ├── onnxt.py             # 后续使用 – ONNX Runtime推理
│   ├── onnx_local.py        # 后续使用 – 本地ONNX推理测试
│   ├── server.py            # 将在第3课中使用（FastAPI后端）
│   └── __init__.py
├── models/
│   ├── resnetv2_50.onnx     # 转换后保存的生成ONNX模型
│   └── imagenet_classes.txt # 将数字预测映射到类别标签
├── data/
│   ├── cat.jpg / dog.jpg    # 用于验证ONNX推理的示例图像
│   ├── cat_base64.txt       # 后续用于Lambda输入负载
│   ├── event.json / payload.json / response.json  # 用于Lambda测试（未来）
├── tests/
│   ├── prepare_event.py     # 准备Lambda测试事件（未来）
│   ├── test_boto3_lambda.py # 测试某中心Lambda调用（未来）
│   └── test_lam.py          # 本地Lambda函数测试（未来）
├── frontend/                # 用于第7–8课（Next.js前端）
├── Dockerfile               # 用于某中心Lambda部署（未来）
├── Dockerfile.local         # 用于本地FastAPI测试（第3课）
├── docker-compose.local.yml # 后续用于多容器设置
├── dev.sh                   # 开发任务的辅助脚本
├── requirements.txt         # 部署的Python依赖项
├── requirements-local.txt   # 本地开发的依赖项
├── .dockerignore
├── .gitignore

快速回顾本课中使用的组件：

  • convert.py：将预训练的ResNetV2-50模型从PyTorch转换为ONNX格式，定义输入张量并保存导出的文件。
  • resnetv2_50.onnx：转换生成的ONNX模型文件，准备好在多个框架和硬件目标上部署。
  • imagenet_classes.txt：一个包含ImageNet类别标签的简单文本文件，用于解释模型的数字输出。
  • cat.jpg / dog.jpg：示例测试图像，用于在继续部署之前验证导出的ONNX模型是否正确执行图像分类。

将ResNetV2-50（TIMM）转换为ONNX

现在开发环境已经设置好，我们将把预训练的ResNetV2-50模型从PyTorch Image Models（TIMM）转换为ONNX格式。ONNX是一种针对不同平台推理进行优化的格式，可确保在某中心Lambda等环境中快速执行。

为什么要转换为ONNX？

  • 跨平台兼容性：ONNX模型可以在PyTorch、TensorFlow和其他框架中使用。
  • 为推理优化：ONNX模型在生产中运行速度比传统PyTorch模型更快。
  • 轻量且高效：ONNX模型占用内存更少，非常适合云部署。

某中心Lambda有执行限制，例如有限的内存和存储，因此使用ONNX Runtime可确保推理得到优化且高效。

将ResNetV2-50（TIMM）导出为ONNX

现在创建一个脚本（convert.py）来：

  1. 从TIMM加载预训练的ResNetV2-50模型。
  2. 定义一个虚拟输入张量（匹配模型的输入大小）。
  3. 转换模型并将其导出为ONNX格式。

创建名为convert.py的新文件并添加以下代码：

import torch
import torch.onnx
import timm
import onnxruntime as ort
import numpy as np
import requests
from PIL import Image
from io import BytesIO

model = timm.create_model('resnetv2_50', pretrained=True)
model = model.eval()
model_script = torch.jit.script(model)

import os

print("exporting onnx...")
# 使用环境变量确定输出路径，默认为模型目录
output_dir = os.getenv("MODEL_PATH", "../models")
output_path = os.path.join(output_dir, "resnetv2_50.onnx")

torch.onnx.export(model_script, torch.randn(1, 3, 224, 224), output_path, verbose=True, input_names=[
                  'input'], output_names=['output'], dynamic_axes={'input': {0: 'batch'}})

代码将预训练的ResNetV2-50模型从PyTorch（通过TIMM）转换为ONNX格式以进行优化推理。首先，加载ResNetV2-50模型并将其设置为评估模式。然后，应用TorchScript（torch.jit）来跟踪模型，确保在将其导出到ONNX之前兼容。torch.onnx.export()函数用于将脚本化模型转换为ONNX，指定形状为（1， 3， 224， 224）的虚拟输入张量来定义输入结构。verbose=True标志记录导出过程，而dynamic_axes允许灵活的批处理大小。最终的ONNX模型（resnetv2_50.onnx）被保存，使其准备好与ONNX Runtime一起进行推理。

理解ONNX模型结构

模型转换后，可以使用onnx检查其结构：

import onnx

# 加载ONNX模型
onnx_model = onnx.load("resnetv2_50.onnx")

# 检查模型结构
print(onnx.helper.printable_graph(onnx_model.graph))

这个可打印的图表示ResNetV2-50模型的ONNX计算图，展示了输入张量、权重和层是如何连接的。它以输入张量%input开始，后面是卷积权重、跨层的批归一化参数和偏置的初始化器。该图定义了模型的阶段和块，包括卷积层（conv）、下采样操作（downsample.conv）和归一化层（norm1）。图的最后部分应用全局平均池化层，然后是全连接（fc）层，并展平输出以生成最终的模型预测。

比较ONNX与PyTorch的AI推理性能

现在我们已经将ResNetv2-50模型转换为ONNX，了解ONNX与PyTorch在推理性能方面的比较非常重要。我们将在以下方面比较两者：

  • 模型文件大小：ONNX是否比PyTorch更轻量？
  • 推理速度：哪种格式推理速度更快？
  • 性能权衡：每种格式的优点和局限性是什么？

比较模型文件大小

对ONNX的最大期望之一是它能减小模型大小，从而更容易部署在某中心Lambda等存储限制常见的环境中。
让我们比较原始PyTorch模型与转换后的ONNX模型的文件大小：

import os

# 保存模型的路径
pytorch_model_path = "resnetv2_50.pth"
onnx_model_path = "resnetv2_50.onnx"

# 检查文件大小（将字节转换为MB）
pytorch_size = os.path.getsize(pytorch_model_path) / (1024 * 1024)
onnx_size = os.path.getsize(onnx_model_path) / (1024 * 1024)

print(f"PyTorch Model Size: {pytorch_size:.2f} MB")
print(f"ONNX Model Size: {onnx_size:.2f} MB")

这里，我们比较了PyTorch（.pth）和ONNX（.onnx）模型的文件大小，以评估ONNX提供的尺寸缩减。我们使用os.path.getsize()获取文件大小（字节），然后将其转换为兆字节（MB）以便于阅读。最后，打印两个模型的大小以确定ONNX是否更轻量，这对于在无服务器环境（如某中心Lambda）中部署很重要，因为存储空间有限。

PyTorch Model Size: 97.74 MB
ONNX Model Size: 97.61 MB

ONNX理想情况下应该比PyTorch小，但正如我们观察到的，两者几乎相同（97MB）。

比较推理速度：PyTorch vs. ONNX

虽然模型大小很重要，但推理速度更为关键，尤其是在延迟影响用户体验的实时AI应用中。
让我们比较使用PyTorch和ONNX Runtime运行推理所需的时间。

PyTorch推理速度

import torch
import timm
import time

# 加载PyTorch模型
pytorch_model = timm.create_model("resnetv2_50", pretrained=True)
pytorch_model.eval()

# 生成虚拟输入张量（批大小：1，通道：3，高度：224，宽度：224）
dummy_input = torch.randn(1, 3, 224, 224)

# 测量推理时间
start_time = time.time()
with torch.no_grad():
    _ = pytorch_model()
pytorch_inference_time = time.time() - start_time

print(f"PyTorch Inference Time: {pytorch_inference_time:.4f} seconds")

接下来，我们测量PyTorch模型处理输入张量并返回输出所需的时间。首先，我们从TIMM加载ResNetV2-50模型并将其设置为评估模式（eval()），禁用梯度更新，因为我们只运行推理。生成一个形状为（1， 3， 224， 224）的随机虚拟输入张量，表示大小为224×224像素的单个RGB图像。然后使用time.time()对模型进行计时以测量推理时间。确保关闭梯度跟踪（torch.no_grad()）以加速执行，因为我们没有训练模型。

PyTorch Inference Time: 0.3639 seconds

PyTorch模型在我们的测试中对单个图像推理耗时0.36秒。

ONNX推理速度

import numpy as np
import time

# 加载ONNX模型
ort_session = ort.InferenceSession("resnetv2_50.onnx")

# 将PyTorch张量转换为NumPy数组
onnx_input = dummy_input.numpy()

# 测量ONNX推理时间
start_time = time.time()
_ = ort_session.run({"input": onnx_input})
onnx_inference_time = time.time()

print(f"ONNX Inference Time: {onnx_inference_time:.4f} seconds")

然后我们使用ONNX Runtime运行推理，并测量它处理与PyTorch中使用的相同输入的速度。它使用onnxruntime.InferenceSession()加载ONNX模型（resnetv2_50.onnx）。由于ONNX需要NumPy格式的输入，因此在推理之前将PyTorch张量转换为NumPy数组（dummy_input.numpy()）。脚本使用time.time()测量所花费的推理时间，与PyTorch版本类似。最后，打印ONNX推理时间以进行比较。

ONNX Inference Time: 0.3017 seconds

ONNX推理比PyTorch的单图像推理快17%。

批处理推理：ONNX表现更出色的地方

虽然单图像推理仅快17%，但我们也测试了批处理推理（batch=128张图像），ONNX的优化在这里真正显示出改进。
批大小 = 128

PyTorch Inference Time: ~40 seconds
ONNX Inference Time: ~20 seconds

在处理较大批次时，ONNX比PyTorch快2倍，这使其成为AI推理工作负载的更好选择。

为什么ONNX对大批次处理更快？

  • ONNX Runtime比PyTorch更有效地优化批处理执行。
  • 开销更低 → PyTorch动态构建每批的计算图，而ONNX使用静态图。
  • ONNX避免了Python的全局解释器锁（GIL），允许多个输入的并行执行。

对于大规模AI部署，ONNX比PyTorch提供了显著的加速改进。

性能权衡和分析

关键要点

  • ONNX推理对单图像推理快17%，对批处理推理（128张图像）快2倍。
  • 模型大小几乎相同（97MB），但轻微的优化可以稍微减小它。
  • ONNX是跨平台的，意味着它可以在不同的AI生态系统中使用。
  • ONNX Runtime更适合大规模AI推理，而PyTorch仍然更适合训练和研究。
    然而，ONNX最适合仅用于推理。如果你需要训练或微调模型，PyTorch仍然是更好的选择。

使用ONNX Runtime测试ONNX模型

现在我们已经确认ONNX对于优化推理是更好的选择，必须使用ONNX Runtime验证ONNX模型的正确性。
我们将通过以下方式测试模型：

  • 在ONNX Runtime上运行推理以确保正确执行。
  • 将ONNX预测与PyTorch进行比较，以确认转换没有影响准确性。
  • 处理潜在的ONNX模型错误并调试任何不一致之处。

在ONNX Runtime上运行推理

我们现在将使用ONNX Runtime运行推理，并验证模型是否正确处理图像。
加载ONNX模型并运行示例推理

import numpy as np
from PIL import Image
import onnxruntime as ort
import torch
import timm
import torchvision.transforms as transforms

# 加载ONNX模型
ort_session = ort.InferenceSession("resnetv2_50.onnx")

# 加载示例图像
image_path = "cat.jpg"  # 更改为任何测试图像
image = Image.open(image_path)

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],               std=[0.229, 0.224, 0.225])
])

# 预处理图像并转换为numpy数组
input_tensor = transform(image).unsqueeze(0).numpy()

# 从模型中获取输入和输出名称
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

# 使用ONNX Runtime运行推理
outputs = ort_session.run([output_name], {input_name: input_tensor})

# 获取预测的类别索引
predicted_class_idx = np.argmax(outputs[0])
print(f"ONNX Model Prediction: Class Index {predicted_class_idx}")

我们使用ONNX Runtime运行推理，以确保转换后的ONNX模型正常运行。我们首先加载ResNetV2-50 ONNX模型和测试图像（例如"cat.jpg"）。然后使用Torchvision变换对图像进行预处理，调整大小为（224， 224），转换为张量，并使用ImageNet的均值和标准差进行归一化。由于ONNX需要NumPy格式的输入，因此在推理之前将处理后的张量转换为NumPy格式。然后模型运行前向传播，并使用argmax()提取预测的类别索引，确定最可能的分类。
ONNX模型应该成功处理图像并输出有效的类别索引。

将ONNX预测与PyTorch进行比较

为了确保ONNX模型在功能上与原始PyTorch模型相同，让我们将同一张图像通过PyTorch版本的模型运行并比较预测结果。

# 加载PyTorch模型
pytorch_model = timm.create_model("resnetv2_50", pretrained=True)
pytorch_model.eval

# 将图像转换为PyTorch张量
input_tensor_torch = transform(image).unsqueeze(0)

# 运行PyTorch推理
with torch.no_grad():
    output_torch = pytorch_model()

# 获取预测的类别索引
predicted_class_torch = torch.argmax(output, dim=-1).item()
print(f"PyTorch Model Prediction: Class Index {predicted_class_torch}")

# 比较ONNX和PyTorch预测
if predicted_class_idx == predicted_torch:
    print("Success! ONNX and PyTorch predictions match.")
else:
    print("Warning: ONNX and PyTorch predictions do not match. Further debugging required.")

然后我们使用PyTorch执行相同的推理。我们从TIMM加载原始的ResNetV2-50模型，并将其设置为评估模式（eval()）以进行推理。同一图像经过相同的预处理，并传递给PyTorch模型。输出是类别上的概率分布，从中选择概率最高的类别。最后，将PyTorch预测与ONNX模型的输出进行比较。
如果ONNX和PyTorch预测匹配，则模型转换成功。如果不匹配，我们可能需要调试ONNX模型转换。

总结

在本课中，我们成功地将PyTorch ResNetV2-50模型转换为ONNX，测试了其性能，并使用ONNX Runtime验证了其正确性。通过分析模型大小和推理速度，我们确认了ONNX在某中心Lambda等无服务器环境的AI推理中具有显著优势。
我们发现ONNX保留了与PyTorch几乎相同的大小（97MB），但我们通过剥离不必要的元数据对其进行了优化。ONNX在单图像推理中快17%，在批处理推理（128张图像）中快2倍。我们验证了ONNX预测与PyTorch匹配，确认转换成功。
随着ONNX模型经过全面测试和优化，我们现在已准备好通过API提供服务并对其进行容器化以进行部署。在下一课中，我们将构建一个基于FastAPI的AI推理服务器，通过API公开我们的ONNX模型。我们还将使用Docker对其进行容器化，确保在部署到某中心Lambda之前具有一致的运行时环境。
更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）或者 我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）

公众号二维码

公众号二维码

posted @ 2026-01-07 16:17  CodeShare  阅读(25)  评论(0)    收藏  举报