当前位置：首页 > news >正文

丹摩 | 基于PyTorch的CIFAR-10图像分类实现

news 2025/6/27 11:00:24

从创建实例开始的新项目流程

第一步：创建实例

登录 DAMODEL 平台。
创建一个 GPU 实例：
- GPU 配置：选择 NVIDIA H800 或其他可用高性能 GPU。
- 系统配置：推荐使用 Ubuntu 20.04，内存 16GB，硬盘 50GB。
- 启动实例后，获取实例的 IP 地址。
- 选择镜像

第二步：连接实例

在这里插入图片描述

登录成功后，你会进入实例的终端界面。

第三步：更新系统和安装基础工具

更新系统：
```
sudo apt update && sudo apt upgrade -y
```

安装 Python 和基础工具：

sudo apt install python3 python3-pip git -y

（可选）安装文本编辑器：
```
sudo apt install vim nano -y
```

第四步：创建项目目录并配置环境

创建项目目录：

mkdir ~/workspace/cifar10_project
cd ~/workspace/cifar10_project

创建并激活虚拟环境：
```
python3 -m venv venv
source venv/bin/activate
```
前面出现venu则表示已经激活虚拟环境了

安装必要的 Python 包：

pip install torch torchvision matplotlib

在这里插入图片描述

第五步：下载数据并初始化项目代码

创建 Python 脚本：
```
vim train_cifar10.py
```

在文件中输入以下代码，加载 CIFAR-10 数据集并定义简单模型：

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim# 数据预处理
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])# 加载 CIFAR-10 数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)# 定义简单卷积神经网络
class SimpleCNN(nn.Module):def __init__(self):super(SimpleCNN, self).__init__()self.conv1 = nn.Conv2d(3, 32, 3, padding=1)self.pool = nn.MaxPool2d(2, 2)self.fc1 = nn.Linear(32 * 16 * 16, 10)def forward(self, x):x = self.pool(torch.relu(self.conv1(x)))x = x.view(-1, 32 * 16 * 16)x = self.fc1(x)return x# 初始化模型、损失函数和优化器
net = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)# 模型训练
for epoch in range(5):  # 训练 5 个周期running_loss = 0.0for inputs, labels in trainloader:optimizer.zero_grad()outputs = net(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()print(f"Epoch {epoch+1}, Loss: {running_loss / len(trainloader)}")print("Finished Training")

保存并退出（按下 Esc，然后输入 :wq）。

第六步：运行训练脚本

运行脚本进行模型训练：

python train_cifar10.py

脚本会下载 CIFAR-10 数据集并训练模型。
训练完成后会输出每个 epoch 的损失值。

第七步：保存和测试模型

保存模型：在脚本末尾添加代码以保存训练好的模型：

torch.save(net.state_dict(), "cifar10_model.pth")
print("Model saved as cifar10_model.pth")

重新运行脚本以保存模型：
```
python train_cifar10.py
```
检查是否生成了 cifar10_model.pth 文件：
```
ls
```

测试模型（可选）：加载保存的模型并在测试集上评估准确率：

net.load_state_dict(torch.load("cifar10_model.pth"))
net.eval()
correct = 0
total = 0
with torch.no_grad():for inputs, labels in testloader:outputs = net(inputs)_, predicted = torch.max(outputs, 1)total += labels.size(0)correct += (predicted == labels).sum().item()print(f"Accuracy on test dataset: {100 * correct / total}%")