PyTorch 模型转换为 ONNX 格式

PyTorch 模型转换为 ONNX 格式

在深度学习领域，模型的可移植性和可解释性是非常重要的。本文将介绍如何使用 PyTorch 训练一个简单的卷积神经网络（CNN）来分类 MNIST 数据集，并将训练好的模型转换为 ONNX 格式。我们还将讨论 PTH 和 ONNX 格式的区别，并介绍如何使用 Netron 可视化 ONNX 模型。

1. PTH 和 ONNX 的区别

PTH 格式

• 定义：PTH 是 PyTorch 框架的专有格式，通常用于保存模型的状态字典（state_dict），包括模型的结构和训练好的参数。

• 兼容性：

  • PTH 文件只能在 PyTorch 中使用，无法直接在 C++ 环境中加载。虽然 PyTorch 提供了 C++ API（LibTorch），但 PTH 文件的加载和使用主要依赖于 Python 环境。
  • 在 C++ 中使用 PTH 文件需要将模型转换为 PyTorch 的 C++ 格式，这可能会增加复杂性和开发时间。

• 用途：

  • PTH 格式适合在 Python 环境中进行模型训练和调试，但在 C++ 中进行模型部署时，通常需要将模型转换为其他格式（如 ONNX）以便于跨平台使用。
  • 在 C++ 中，使用 PTH 文件的灵活性较低，尤其是在需要与其他框架或系统集成时。

ONNX 格式

• 定义：ONNX（Open Neural Network Exchange）是一个开放的深度学习模型交换格式，旨在促进不同深度学习框架之间的互操作性。

• 兼容性：

  • ONNX 文件可以在多个深度学习框架中使用，包括 PyTorch、TensorFlow、Caffe2 等，这使得它在 C++ 环境中的兼容性更强。
  • ONNX 模型可以通过 ONNX Runtime、TensorRT、OpenVINO 等推理引擎在 C++ 中高效运行，支持多种硬件加速。

• 用途：

  • ONNX 格式非常适合模型的部署和推理，特别是在需要跨平台或跨框架使用时。它允许开发者在 C++ 中轻松加载和运行模型，而无需依赖于 Python 环境。
  • 在 C++ 中，使用 ONNX 模型可以简化工程化流程，便于与其他系统集成，提升模型的可移植性和可扩展性。

总结

在 C++ 进行深度学习模型的工程化时，选择 ONNX 格式通常更为合适，因为它提供了更好的跨平台兼容性和灵活性。PTH 格式虽然在 PyTorch 环境中非常方便，但在 C++ 中的使用受到限制，通常需要额外的转换步骤。ONNX 的开放性和广泛支持使其成为在多种环境中部署深度学习模型的首选格式。

2. 训练 MNIST 数据集的 CNN 模型

以下是使用 PyTorch 训练 MNIST 数据集的完整代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision import datasets
from torch.utils.data import DataLoader

# 检查是否支持 MPS
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print(f"Using device: {device}")

# 1. 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # MNIST 数据集的均值和标准差
])

# 下载 MNIST 数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

# 2. 定义 CNN 模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)  # 输入通道为1，输出通道为32
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)  # 输入通道为32，输出通道为64
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)  # 最大池化层
        self.fc1 = nn.Linear(64 * 7 * 7, 128)  # 全连接层
        self.fc2 = nn.Linear(128, 10)  # 输出层

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))  # 第一层卷积 + 激活 + 池化
        x = self.pool(torch.relu(self.conv2(x)))  # 第二层卷积 + 激活 + 池化
        x = x.view(x.size(0), -1)  # 展平输入
        x = torch.relu(self.fc1(x))  # 第一个全连接层
        x = self.fc2(x)  # 输出层
        return x

# 3. 训练模型
model = SimpleCNN().to(device)  # 将模型移动到 MPS 设备
criterion = nn.CrossEntropyLoss()  # 损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 优化器

# 训练过程
num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)  # 将数据移动到 MPS 设备
        optimizer.zero_grad()  # 清空梯度
        outputs = model(images)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

# 4. 评估模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        images, labels = images.to(device), labels.to(device)  # 将数据移动到 MPS 设备
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)  # 获取预测结果
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the model on the test images: {100 * correct / total:.2f}%')

# 5. 转换为 ONNX 格式
onnx_file_path = 'mnist_cnn_model.onnx'
dummy_input = torch.randn(1, 1, 28, 28).to(device)  # 示例输入，形状为 [batch_size, channels, height, width]
torch.onnx.export(model, dummy_input, onnx_file_path, export_params=True,
                  opset_version=11, do_constant_folding=True,
                  input_names=['input'], output_names=['output'],
                  dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}})

print(f'Model has been converted to ONNX format and saved as {onnx_file_path}.')

3. 使用 Netron 可视化 ONNX 模型

一旦您将模型转换为 ONNX 格式，您可以使用 Netron 来可视化模型结构。Netron 是一个开源的模型可视化工具，支持多种深度学习框架的模型文件格式，包括 ONNX。

使用步骤：

1. 下载 Netron：

   • 您可以访问 Netron 的官方网站 在线使用，或者下载桌面版本。

2. 打开 ONNX 模型：

   • 如果使用在线版本，直接将 mnist_cnn_model.onnx 文件拖放到浏览器窗口中。
   • 如果使用桌面版本，打开 Netron 应用，选择“File” > “Open Model”，然后选择您的 ONNX 文件。

3. 查看模型结构：

   • 在 Netron 中，您可以查看模型的层次结构、输入输出形状、参数数量等信息。通过可视化，您可以更好地理解模型的设计和工作原理。