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
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
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
def compile_model(net):
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
return criterion, optimizer
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.0
print('训练完成')
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 = data
images, 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]] += 1
total_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 = f
print(f"当前使用设备: {device}")
train(net, trainloader, criterion, optimizer, device)
test(net, testloader, classes, device)
sys.stdout = original_stdout # 恢复控制台输出
print("训练完成！结果已保存到 cifar10_result.txt ")
