CNN 算法可视化演示

完整 CNN 实现 (PyTorch)

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 定义CNN模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一层卷积：输入通道1，输出通道32，卷积核大小3x3
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        # 第二层卷积：输入通道32，输出通道64，卷积核大小3x3
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        # 最大池化层：窗口大小2x2
        self.pool1 = nn.MaxPool2d(2)
        # Dropout层：随机失活比例0.25
        self.dropout1 = nn.Dropout2d(0.25)
        # 自适应平均池化层：输出大小为1x1
        self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1))
        # 全连接层1：输入256维，输出10维
        self.fc1 = nn.Linear(256, 10)
        # Dropout层：随机失活比例0.5
        self.dropout2 = nn.Dropout(0.5)
    
    def forward(self, x):
        # 卷积 -> 激活 -> 卷积 -> 激活 -> 池化 -> Dropout
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = self.pool1(x)
        x = self.dropout1(x)
        # 应用自适应平均池化
        x = self.adaptive_pool(x)
        # 展平张量
        x = torch.flatten(x, 1)
        # 全连接层 -> Dropout -> 全连接层
        x = torch.relu(self.fc1(x))
        x = self.dropout2(x)
        # 使用log_softmax激活函数
        return torch.log_softmax(x, dim=1)

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 实例化模型并移动到设备
model = CNN().to(device)

# 定义超参数
batch_size = 64
epochs = 10
learning_rate = 0.001

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 下载训练和测试数据集
torchvision.datasets.MNIST('data', train=True, download=True, transform=transform)
dataset = torchvision.datasets.MNIST('data', train=False, download=True, transform=transform)

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('data', train=True, download=False, transform=transform),
    batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('data', train=False, download=False, transform=transform),
    batch_size=batch_size, shuffle=True)

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

# 训练函数
def train(model, device, train_loader, optimizer, criterion, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')

# 测试函数
def test(model, device, test_loader, criterion):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    
    test_loss /= len(test_loader.dataset)
    print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n')

# 训练循环
for epoch in range(1, epochs + 1):
    train(model, device, train_loader, optimizer, criterion, epoch)
    test(model, device, test_loader, criterion)