本文将介绍如何使用 PyTorch 训练一个简单的卷积神经网络（CNN）模型，并将训练好的模型导出为 ONNX 格式，之后使用 ONNX Runtime 在 Python 中加载并进行推理。
在开始使用模型前，需要安装以下 Python 库。为了提高下载速度，我们将使用国内镜像源进行安装。以下是各个库的安装步骤。
1. 安装 onnxruntime
onnxruntime 是用于加载和运行 ONNX 模型的库。可以通过以下命令使用国内镜像源进行安装：
pip install onnxruntime -i https://pypi.tuna.tsinghua.edu.cn/simple
2. 安装 pandas
pandas 是用于数据处理和分析的库，特别适合处理表格数据。通过以下命令安装：
pip install pandas -i https://pypi.tuna.tsinghua.edu.cn/simple
3. 安装 Pillow（PIL）
Pillow 是 Python 的图像处理库，PIL 是它的别名。用于图像读取和处理。安装命令如下：
pip install Pillow -i https://pypi.tuna.tsinghua.edu.cn/simple
4. 安装 scikit-learn
scikit-learn 是常用的机器学习库，包含数据预处理、模型训练等功能。使用以下命令安装：
pip install scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple
1. 数据集和数据加载
我们将使用一个简单的手写数字分类数据集——Pendigits，这是一个包含手写数字的 CSV 文件。首先，我们需要处理数据并创建一个 PyTorch 数据集。
自定义数据集类：
import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from sklearn.preprocessing import StandardScaler
class PendigitsDataset(Dataset):
def __init__(self, csv_file):
self.data = pd.read_csv(csv_file, header=0)
self.features = self.data.iloc[:, 1:-1].values
self.labels = self.data.iloc[:, -1].values
# 数据标准化
self.scaler = StandardScaler()
self.features = self.scaler.fit_transform(self.features)
def __len__(self):
return len(self.features)
def __getitem__(self, idx):
sample = torch.tensor(self.features[idx], dtype=torch.float32)
label = torch.tensor(self.labels[idx], dtype=torch.long)
sample = sample.view(1, 28, 28) # 转换为28x28图像格式
return sample, label
在 PendigitsDataset 类中，__init__ 方法加载 CSV 文件并进行数据标准化；__getitem__ 方法将每个样本转换为一个 28x28 的图像格式。
数据加载器：
trainset = PendigitsDataset(csv_file='path_to_csv_file')
trainloader = DataLoader(trainset, batch_size=32, shuffle=True)
2. 定义卷积神经网络模型（CNN）
我们使用两个卷积层和两个全连接层来构建一个简单的 CNN。
import torch.nn as nn
class CNN(nn.Module):
def __init__(self, output_size=10):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, output_size)
self.relu = nn.ReLU()
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
x = self.relu(self.conv1(x))
x = self.pool(x)
x = self.relu(self.conv2(x))
x = self.pool(x)
x = x.view(-1, 64 * 7 * 7) # 展平
x = self.relu(self.fc1(x))
x = self.fc2(x)
return self.softmax(x)
3. 训练模型
接下来，我们将训练模型并使用优化器和损失函数进行训练。
import torch.optim as optim
# 初始化模型，损失函数，优化器
model = CNN(output_size=10)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
epochs = 100
for epoch in range(epochs):
running_loss = 0.0
for inputs, labels in trainloader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}, Loss: {running_loss/len(trainloader)}")
4. 导出模型为 ONNX 格式
完成训练后，我们可以将模型导出为 ONNX 格式，以便在其他平台（如 ONNX Runtime）中进行推理。
import torch.onnx
# 导出模型为ONNX格式
onnx_path = 'cnn_model.onnx'
dummy_input = torch.randn(1, 1, 28, 28).to(device) # 创建一个与训练数据形状相同的虚拟输入
torch.onnx.export(model, dummy_input, onnx_path, input_names=['input'], output_names=['output'])
print(f"模型已导出为 ONNX 格式，路径为 {onnx_path}")
完整代码
import torch
import torch.onnx
import torch.nn as nn
import torch.optim as optim
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
# 自定义数据集类
class PendigitsDataset(Dataset):
def __init__(self, csv_file, transform=None):
# 读取CSV文件，并跳过第一行（列名）
self.data = pd.read_csv(csv_file, header=0) # header=0 表示第一行是列名
self.transform = transform
# 去除不需要训练的列（如 a1, q1）
# 假设这些列在前几列，并且标签列是最后一列
# 获取特征数据（所有列，去掉最后一列标签）
self.features = self.data.iloc[:, 1:-1].values # 去掉第一列和最后一列（即标签）
self.labels = self.data.iloc[:, -1].values # 标签列
# 数据标准化（如果需要）
self.scaler = StandardScaler()
self.features = self.scaler.fit_transform(self.features) # 标准化特征数据
def __len__(self):
return len(self.features)
def __getitem__(self, idx):
sample = torch.tensor(self.features[idx], dtype=torch.float32)
label = torch.tensor(self.labels[idx], dtype=torch.long)
if self.transform:
sample = self.transform(sample)
return sample, label
# 定义MLP模型
class MLP(nn.Module):
def __init__(self, input_size=16, hidden_size=128, output_size=10):
super(MLP, self).__init__()
self.fc1 = nn.Linear(input_size, hidden_size)
self.fc2 = nn.Linear(hidden_size, hidden_size)
self.fc3 = nn.Linear(hidden_size, output_size)
self.relu = nn.ReLU()
self.softmax = nn.Softmax(dim=1)
def forward(self, x):
x = self.relu(self.fc1(x))
x = self.relu(self.fc2(x))
x = self.fc3(x)
return self.softmax(x)
# 数据预处理
transform = None # 你可以根据需要定义transform
# 使用Pendigits数据集
trainset = PendigitsDataset(csv_file=r'pendigits_txt.csv', transform=transform)# 你的csv格式数据集路径
trainloader = DataLoader(trainset, batch_size=32, shuffle=True) # 增加batch_size
# 初始化模型，损失函数，优化器
model = MLP(input_size=trainset.features.shape[1]) # 输入大小是特征的数量
# 使用GPU训练（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 加载模型（如果有保存的模型）
checkpoint_path = 'mlp_model_checkpoint.pth'
start_epoch = 0
try:
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1 # 从保存的epoch继续
print(f"模型从epoch {start_epoch}开始恢复训练")
except FileNotFoundError:
print("没有找到之前的模型，开始新训练")
# 训练模型
epochs = 100
for epoch in range(start_epoch, start_epoch + epochs):
running_loss = 0.0
for inputs, labels in trainloader:
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}, Loss: {running_loss/len(trainloader)}")
# 每训练一定的epoch保存一次模型
if (epoch + 1) % 20 == 0: # 每20个epoch保存一次
checkpoint_filename = f'mlp_model_checkpoint_{epoch-19}-{epoch}.pth'
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
}, checkpoint_filename)
print(f"模型已保存至 {checkpoint_filename}")
# 导出ONNX模型
dummy_input = torch.randn(1, trainset.features.shape[1]).to(device) # 确保dummy_input也在GPU上
onnx_path = "mlp_model.onnx"
torch.onnx.export(model, dummy_input, onnx_path, input_names=['input'], output_names=['output'])
print(f"模型已导出到 {onnx_path}")
5. 使用 ONNX Runtime 进行推理
我们可以使用 onnxruntime 来加载训练好的 ONNX 模型，并使用它进行推理。
import onnx
import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms
# 加载训练好的 ONNX 模型
model_path = "cnn_model.onnx"
session = ort.InferenceSession(model_path)
# 预处理函数：将图片调整为 28x28 尺寸并转为灰度图像、张量
transform = transforms.Compose([
transforms.Grayscale(num_output_channels=1), # 转为灰度图
transforms.Resize((28, 28)), # 调整为 28x28 尺寸
transforms.ToTensor(), # 转为张量
transforms.Normalize((0.5,), (0.5,)) # 归一化
])
# 加载并预处理输入图片
image_path = "path_to_image.jpg" # 替换为你的图片路径
image = Image.open(image_path)
image = transform(image).numpy() # 转为 NumPy 数组
# 进行推理（预测）
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
outputs = session.run([output_name], {input_name: image.astype(np.float32)})
# 获取预测结果
predicted_class = np.argmax(outputs[0], axis=1)
print(f"预测的类是: {predicted_class[0]}")
6. 总结
本文介绍了如何在 PyTorch 中训练一个 CNN 模型，并将其导出为 ONNX 格式。我们还展示了如何使用 ONNX Runtime 在 Python 中加载导出的模型并进行推理。通过这种方式，你可以轻松地将 PyTorch 模型部署到不同的平台并进行推理。
扫一扫
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
