p66

点击查看代码

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 1. 加载和预处理数据
def load_and_preprocess_data():
    """
    加载和预处理CIFAR-10数据集
    """
    # 数据预处理
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    # 加载训练集和测试集
    trainset = torchvision.datasets.CIFAR10(
        root='./data', train=True, download=True, transform=transform_train)
    
    testset = torchvision.datasets.CIFAR10(
        root='./data', train=False, download=True, transform=transform_test)

    # 创建数据加载器
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=128, shuffle=True, num_workers=2)
    
    testloader = torch.utils.data.DataLoader(
        testset, batch_size=100, shuffle=False, num_workers=2)

    # CIFAR-10类别名称
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 
               'dog', 'frog', 'horse', 'ship', 'truck')

    return trainloader, testloader, classes

# 2. 构建网络
class CIFAR10Net(nn.Module):
    """
    构建CIFAR-10图像分类网络
    """
    def __init__(self):
        super(CIFAR10Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        
        self.pool = nn.MaxPool2d(2, 2)
        self.dropout = nn.Dropout(0.25)
        
        self.fc1 = nn.Linear(128 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        # 卷积层1
        x = self.pool(F.relu(self.conv1(x)))
        x = self.dropout(x)
        
        # 卷积层2
        x = self.pool(F.relu(self.conv2(x)))
        x = self.dropout(x)
        
        # 卷积层3
        x = self.pool(F.relu(self.conv3(x)))
        x = self.dropout(x)
        
        # 展平
        x = x.view(-1, 128 * 4 * 4)
        
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        
        return x

def build_network():
    """
    构建网络模型
    """
    model = CIFAR10Net().to(device)
    return model

# 3. 编译网络
def compile_network(model):
    """
    编译网络：设置优化器和损失函数
    """
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
    
    return criterion, optimizer, scheduler

# 4. 训练网络
def train_network(model, trainloader, criterion, optimizer, scheduler, num_epochs=50):
    """
    训练网络模型
    """
    train_losses = []
    model.train()
    
    for epoch in range(num_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 = model(inputs)
            loss = criterion(outputs, labels)
            
            # 反向传播
            loss.backward()
            optimizer.step()
            
            # 统计损失
            running_loss += loss.item()
            
            # 每100个batch打印一次
            if i % 100 == 99:
                print(f'Epoch: {epoch+1} [{i+1}/{len(trainloader)} ({100.*(i+1)/len(trainloader):.0f}%)] '
                      f'Loss: {running_loss/100:.6f}')
                train_losses.append(running_loss/100)
                running_loss = 0.0
        
        # 更新学习率
        scheduler.step()
        
        # 每个epoch结束后打印信息
        print(f'Epoch {epoch+1} completed')
    
    return train_losses

# 评估模型性能
def evaluate_model(model, testloader, classes):
    """
    评估模型在测试集上的性能
    """
    model.eval()
    correct = 0
    total = 0
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))
    
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            
            # 计算每个类别的准确率
            c = (predicted == labels).squeeze()
            for i in range(len(labels)):
                label = labels[i]
                class_correct