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

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

article 2024/11/28 10:34:38

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

第一步：创建实例

登录 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.0
    for 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}%")