import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
===================== 1. 数据加载与预处理 =====================
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
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=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
classes = ('飞机', '汽车', '鸟', '猫', '鹿', '狗', '青蛙', '马', '船', '卡车')
===================== 2. 构建卷积神经网络 =====================
class CIFAR10Net(nn.Module):
def init(self):
super(CIFAR10Net, self).init()
self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(64, 128, 3, padding=1)self.relu2 = nn.ReLU()self.pool2 = nn.MaxPool2d(2, 2)self.fc1 = nn.Linear(128 * 8 * 8, 512)self.relu3 = nn.ReLU()self.fc2 = nn.Linear(512, 10)def forward(self, x):x = self.pool1(self.relu1(self.conv1(x)))x = self.pool2(self.relu2(self.conv2(x)))x = x.view(-1, 128 * 8 * 8)x = self.relu3(self.fc1(x))x = self.fc2(x)return x
model = CIFAR10Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
===================== 3. 训练网络 =====================
epochs = 10
train_losses = []
train_accs = []
for epoch in range(epochs):
running_loss = 0.0
correct = 0
total = 0
for inputs, labels in train_loader:optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()_, predicted = outputs.max(1)total += labels.size(0)correct += predicted.eq(labels).sum().item()epoch_loss = running_loss / len(train_loader)
epoch_acc = 100. * correct / total
train_losses.append(epoch_loss)
train_accs.append(epoch_acc)print(f'Epoch {epoch+1}/{epochs} | Loss: {epoch_loss:.3f} | Accuracy: {epoch_acc:.2f}%')
print("训练完成！")
===================== 4. 模型评估（测试集） =====================
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in test_loader:
outputs = model(inputs)
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
test_acc = 100. * correct / total
print(f'测试集准确率: {test_acc:.2f}%')
===================== （可选）可视化训练过程 =====================
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs+1), train_losses)
plt.title('训练损失')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.subplot(1, 2, 2)
plt.plot(range(1, epochs+1), train_accs)
plt.title('训练准确率')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.tight_layout()
plt.show()