知识点回顾：
- 彩色和灰度图片测试和训练的规范写法：封装在函数中
- 展平操作：除第一个维度batchsize外全部展平
- dropout操作：训练阶段随机丢弃神经元，测试阶段eval模式关闭dropout
作业：仔细学习下测试和训练代码的逻辑，这是基础，这个代码框架后续会一直沿用，后续的重点慢慢就是转向模型定义阶段了。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
# 设备配置
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# 数据预处理函数 - 处理灰度图像(MNIST)
def get_mnist_loaders(batch_size=64):
# 灰度图像归一化
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)) # MNIST数据集的均值和标准差
])
train_dataset = datasets.MNIST(
root='./data',
train=True,
download=True,
transform=transform
)
test_dataset = datasets.MNIST(
root='./data',
train=False,
download=True,
transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
return train_loader, test_loader
# 数据预处理函数 - 处理彩色图像(CIFAR-10)
def get_cifar10_loaders(batch_size=64):
# 彩色图像归一化
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # CIFAR-10的均值和标准差
])
train_dataset = datasets.CIFAR10(
root='./data',
train=True,
download=True,
transform=transform
)
test_dataset = datasets.CIFAR10(
root='./data',
train=False,
download=True,
transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
return train_loader, test_loader
# 定义包含展平和dropout的CNN模型
class SimpleCNN(nn.Module):
def __init__(self, in_channels=1, num_classes=10):
super(SimpleCNN, self).__init__()
# 卷积层部分
self.conv_layers = nn.Sequential(
nn.Conv2d(in_channels, 32, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2),
nn.Dropout(0.25), # 第一个dropout层
nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2),
nn.Dropout(0.25) # 第二个dropout层
)
# 展平操作后接全连接层
self.fc_layers = nn.Sequential(
nn.Flatten(), # 展平操作，保持batch维度不变
nn.Linear(64 * 7 * 7, 128), # 假设输入尺寸为32x32，经过两次池化后为8x8，64通道
nn.ReLU(),
nn.Dropout(0.5), # 全连接层后的dropout
nn.Linear(128, num_classes)
)
def forward(self, x):
x = self.conv_layers(x)
x = self.fc_layers(x)
return x
# 训练函数 - 规范的训练流程
def train(model, train_loader, criterion, optimizer, epoch, device):
model.train() # 切换到训练模式，启用dropout
running_loss = 0.0
correct = 0
total = 0
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()
# 统计训练信息
running_loss += loss.item()
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
if batch_idx % 100 == 0:
print(f'Epoch: {epoch} | Batch: {batch_idx} | Loss: {loss.item():.4f} | Acc: {100.*correct/total:.2f}%')
epoch_loss = running_loss / len(train_loader)
epoch_acc = 100. * correct / total
return epoch_loss, epoch_acc
# 测试函数 - 规范的测试流程
def test(model, test_loader, criterion, device):
model.eval() # 切换到测试模式，关闭dropout
test_loss = 0
correct = 0
total = 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()
# 计算准确率
_, predicted = output.max(1)
total += target.size(0)
correct += predicted.eq(target).sum().item()
test_loss /= len(test_loader)
test_acc = 100. * correct / total
print(f'Test Loss: {test_loss:.4f} | Test Accuracy: {test_acc:.2f}%')
return test_loss, test_acc
# 主函数 - 整合整个流程
def main(use_color=False):
# 选择数据集
if use_color:
print("Using CIFAR-10 (color images) dataset...")
train_loader, test_loader = get_cifar10_loaders(batch_size=128)
in_channels = 3 # 彩色图像3通道
else:
print("Using MNIST (grayscale images) dataset...")
train_loader, test_loader = get_mnist_loaders(batch_size=128)
in_channels = 1 # 灰度图像1通道
# 初始化模型
model = SimpleCNN(in_channels=in_channels, num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练和测试循环
num_epochs = 5
train_losses, train_accs, test_losses, test_accs = [], [], [], []
for epoch in range(1, num_epochs + 1):
print(f"\nEpoch {epoch}/{num_epochs}")
train_loss, train_acc = train(model, train_loader, criterion, optimizer, epoch, device)
test_loss, test_acc = test(model, test_loader, criterion, device)
train_losses.append(train_loss)
train_accs.append(train_acc)
test_losses.append(test_loss)
test_accs.append(test_acc)
# 绘制训练过程
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.title('Loss Curve')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.subplot(1, 2, 2)
plt.plot(train_accs, label='Train Accuracy')
plt.plot(test_accs, label='Test Accuracy')
plt.title('Accuracy Curve')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.legend()
plt.tight_layout()
plt.show()
if __name__ == "__main__":
# 运行灰度图像版本(MNIST)
main(use_color=False)
# 取消注释以下行运行彩色图像版本(CIFAR-10)
# main(use_color=True)