import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from multiprocessing import freeze_support
import sys
1. 加载和预处理数据
def load_data():
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,download=True,transform=transform
)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,shuffle=True,num_workers=2
)testset = torchvision.datasets.CIFAR10(root='./data', train=False,download=True,transform=transform
)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,shuffle=False,num_workers=2
)classes = ('plane', 'car', 'bird', 'cat','deer', 'dog', 'frog', 'horse', 'ship', 'truck')return trainloader, testloader, classes
2. 构建网络
class Net(nn.Module):
def init(self):
super().init()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):x = self.pool(torch.relu(self.conv1(x)))x = self.pool(torch.relu(self.conv2(x)))x = torch.flatten(x, 1) # 展平x = torch.relu(self.fc1(x))x = torch.relu(self.fc2(x))x = self.fc3(x)return x
3. 编译网络（定义损失函数和优化器）
def compile_model(net):
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
return criterion, optimizer
4. 训练网络（已同步设备）
def train(net, trainloader, criterion, optimizer, device, epochs=2):
for epoch in range(epochs):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
# 核心：数据与模型设备同步
inputs, labels = inputs.to(device), labels.to(device)
# 梯度清零optimizer.zero_grad()# 前向计算 + 反向传播 + 优化参数outputs = net(inputs)loss = criterion(outputs, labels)loss.backward()optimizer.step()# 打印训练日志running_loss += loss.item()if i % 2000 == 1999: # 每2000个batch打印一次print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')running_loss = 0.0print('训练完成')
5. 测试网络（已同步设备）
def test(net, testloader, classes, device):
correct = 0
total = 0
# 测试时不计算梯度，加快速度
with torch.no_grad():
for data in testloader:
images, labels = data
# 数据与模型设备同步
images, labels = images.to(device), labels.to(device)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'测试集整体准确率: {100 * correct // total} %')# 按类别统计准确率
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}with torch.no_grad():for data in testloader:images, labels = dataimages, labels = images.to(device), labels.to(device)outputs = net(images)_, predictions = torch.max(outputs, 1)# 统计每个类别的预测结果for label, prediction in zip(labels, predictions):if label == prediction:correct_pred[classes[label]] += 1total_pred[classes[label]] += 1# 打印各类别准确率
for classname, correct_count in correct_pred.items():accuracy = 100 * float(correct_count) / total_pred[classname]print(f'类别: {classname:5s} 准确率: {accuracy:.1f} %')
if name == 'main':
freeze_support() # 解决Windows多进程问题
# 自动选择设备（有GPU用GPU，无则用CPU）
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"当前使用设备: {device}")
# 加载数据、初始化模型和优化器
trainloader, testloader, classes = load_data()
net = Net().to(device) # 模型放到指定设备
criterion, optimizer = compile_model(net)# 重定向输出到文件，同时保留控制台打印
original_stdout = sys.stdout
with open('cifar10_result.txt', 'w') as f:sys.stdout = fprint(f"当前使用设备: {device}")train(net, trainloader, criterion, optimizer, device)test(net, testloader, classes, device)sys.stdout = original_stdout # 恢复控制台输出print("训练完成！结果已保存到 cifar10_result.txt ")
