运行代码
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor, Normalize
import time
1. 数据加载（简化预处理）
transform = Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
train_dataset = CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)
2. 简化模型结构
class SimpleModel(nn.Module):
def init(self):
super().init()
self.conv = nn.Conv2d(3, 16, kernel_size=3, padding=1)
self.relu = nn.ReLU()
self.pool = nn.MaxPool2d(2, 2)
self.flat = nn.Flatten()
self.fc = nn.Linear(16 * 16 * 16, 10) # 32/2/2=8？不对，32经过两次池化（22）后是8×8？哦，我之前算错了，重新来：32→池化后16→再池化后8，所以168*8=1024。抱歉之前的错误，这里修正。
self.fc = nn.Linear(16 * 8 * 8, 10)
def forward(self, x):x = self.relu(self.conv(x))x = self.pool(x)x = self.relu(self.conv(x)) # 再加一层卷积x = self.pool(x)x = self.flat(x)x = self.fc(x)return x
model = SimpleModel()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
3. 损失函数与优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
4. 简化训练（减少轮数）
epochs = 10
start_time = time.time()
for epoch in range(epochs):
model.train()
running_loss = 0.0
for inputs, targets in train_loader:
inputs, targets = inputs.to(device), targets.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f"Epoch {epoch+1}, 训练损失: {running_loss/len(train_loader):.3f}")
print(f"训练耗时: {time.time()-start_time:.2f} 秒")
5. 简单评估
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, targets in test_loader:
inputs, targets = inputs.to(device), targets.to(device)
outputs = model(inputs)
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
print(f"测试准确率: {correct/total:.3f}")
测试结果图片
