使用 CNN 进行图像分类

在本教程中，我们将构建一个简单的卷积神经网络 (CNN) 来对 MNIST 数据集中的手写数字进行分类。

1. 导入库

import tensorplay as tp
import tensorplay.nn as nn
import tensorplay.optim as optim
from tensorplay.utils.data import DataLoader
from tensorplay.datasets import MNIST

2. 定义模型

我们将使用两个卷积层，后面跟着两个全连接层。

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.flatten(1)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = CNN()

3. 训练设置

# 超参数
batch_size = 64
learning_rate = 0.001
epochs = 5

# 数据加载器
train_dataset = MNIST(root='./data', train=True, download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

4. 训练循环

for epoch in range(epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 100 == 0:
            print(f'Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')

5. 评估

model.eval()
with tp.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = tp.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print(f'Test Accuracy: {100 * correct / total}%')