pytorch66页实验题

点击查看代码

import torch
from torch import optim, nn
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
import numpy as np
from sklearn.metrics import classification_report
from PIL import Image
import time
from matplotlib import pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据加载和预处理
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)),
])

# 加载CIFAR-10数据集
def load_data():
    train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
    test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)

    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2, pin_memory=True)
    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2, pin_memory=True)

    return train_loader, test_loader, test_dataset

# 定义改进的MYVGG模型
class MYVGG(nn.Module):
    def __init__(self, num_classes=10):
        super(MYVGG, self).__init__()
        self.features = nn.Sequential(
            # Block 1
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),

            # Block 2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),

            # Block 3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(256 * 4 * 4, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

# 训练函数
def train(model, train_loader, test_loader, criterion, optimizer, scheduler, epoch_num=50):
    model.train()
    train_loss = []
    train_acc = []
    test_acc = []
    
    best_acc = 0.0
    
    for epoch in range(epoch_num):
        start_time = time.time()
        running_loss = 0.0
        correct = 0
        total = 0
        
        # 训练阶段
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)

            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

        # 更新学习率
        scheduler.step()
        
        epoch_loss = running_loss / len(train_loader)
        epoch_acc = 100.0 * correct / total
        train_loss.append(epoch_loss)
        train_acc.append(epoch_acc)

        # 测试阶段
        current_test_acc = test(model, test_loader, verbose=False)
        test_acc.append(current_test_acc)
        
        # 保存最佳模型
        if current_test_acc > best_acc:
            best_acc = current_test_acc
            torch.save(model.state_dict(), 'best_cifar10_model.pth')

        end_time = time.time()
        current_lr = optimizer.param_groups[0]['lr']
        print(f"Epoch [{epoch+1}/{epoch_num}], Loss: {epoch_loss:.4f}, "
              f"Train Acc: {epoch_acc:.2f}%, Test Acc: {current_test_acc:.2f}%, "
              f"LR: {current_lr:.6f}, Time: {end_time-start_time:.2f}s")
    
    # 加载最佳模型
    model.load_state_dict(torch.load('best_cifar10_model.pth'))
    print(f"最佳测试准确率: {best_acc:.2f}%")
    
    return train_loss, train_acc, test_acc

# 测试函数
def test(model, test_loader, verbose=True):
    model.eval()
    all_pred = []
    all_label = []

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            outputs = model(data)
            _, predicted = torch.max(outputs.data, 1)
            all_pred.extend(predicted.cpu().numpy())
            all_label.extend(target.cpu().numpy())
        
        all_pred = np.array(all_pred)
        all_label = np.array(all_label)
        accuracy = (all_pred == all_label).mean()
        accuracy = 100.0 * accuracy
        
        if verbose:
            print(f'测试准确率: {accuracy:.2f}%')
            
            print("分类效果评估：")
            class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                          'dog', 'frog', 'horse', 'ship', 'truck']
            report = classification_report(all_label, all_pred, target_names=class_names, digits=4)
            print(report)
        
        return accuracy

# 可视化训练过程
def plot_results(train_loss, train_acc, test_acc, epoch_num):
    plt.figure(figsize=(15, 5))
    
    plt.subplot(1, 3, 1)
    plt.plot(range(1, epoch_num+1), train_loss)
    plt.title("Training Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.grid(True)
    
    plt.subplot(1, 3, 2)
    plt.plot(range(1, epoch_num+1), train_acc, label='Train Accuracy')
    plt.plot(range(1, epoch_num+1), test_acc, label='Test Accuracy')
    plt.title("Training and Test Accuracy")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy (%)")
    plt.legend()
    plt.grid(True)
    
    plt.subplot(1, 3, 3)
    plt.plot(range(1, epoch_num+1), train_acc, 'b-', label='Train Acc')
    plt.plot(range(1, epoch_num+1), test_acc, 'r-', label='Test Acc')
    plt.plot(range(1, epoch_num+1), train_loss, 'g--', label='Train Loss (scaled)')
    # 缩放损失以便在同一图表中显示
    scaled_loss = [loss * 100 for loss in train_loss]  # 放大100倍以便观察
    plt.plot(range(1, epoch_num+1), scaled_loss, 'g--', alpha=0.7, label='Train Loss (scaled)')
    plt.title("综合训练曲线")
    plt.xlabel("Epoch")
    plt.ylabel("Value")
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig('training_results.png', dpi=300, bbox_inches='tight')
    plt.show()

# 显示一些测试样本的预测结果
def show_predictions(model, test_dataset, num_samples=12):
    model.eval()
    indices = np.random.choice(len(test_dataset), num_samples)
    
    plt.figure(figsize=(15, 10))
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                  'dog', 'frog', 'horse', 'ship', 'truck']
    
    with torch.no_grad():
        for i, idx in enumerate(indices):
            image, label = test_dataset[idx]
            image_input = image.unsqueeze(0).to(device)
            output = model(image_input)
            _, predicted = torch.max(output, 1)
            
            plt.subplot(3, 4, i+1)
            # 反归一化以便显示
            image = image * torch.tensor([0.2023, 0.1994, 0.2010]).view(3, 1, 1) + torch.tensor([0.4914, 0.4822, 0.4465]).view(3, 1, 1)
            image = torch.clamp(image, 0, 1)
            plt.imshow(image.permute(1, 2, 0).cpu().numpy())
            
            true_label = class_names[label]
            pred_label = class_names[predicted.item()]
            color = 'green' if predicted.item() == label else 'red'
            
            plt.title(f'True: {true_label}\nPred: {pred_label}', color=color)
            plt.axis('off')
    
    plt.tight_layout()
    plt.savefig('sample_predictions.png', dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == '__main__':
    print(f"24信计2班 王晶莹 2024310143126")
    print(f"Device: {device}")
    
    epoch_num = 50
    
    # 加载数据
    train_loader, test_loader, test_dataset = load_data()
    
    # 初始化模型
    model = MYVGG().to(device)
    
    # 定义损失函数和优化器
    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)
    
    # 打印模型参数量
    total_params = sum(p.numel() for p in model.parameters())
    print(f"模型总参数量: {total_params:,}")
    
    # 训练模型
    train_loss, train_acc, test_acc = train(model, train_loader, test_loader, 
                                          criterion, optimizer, scheduler, epoch_num)
    
    # 最终测试
    print("\n最终测试结果:")
    final_accuracy = test(model, test_loader, verbose=True)
    
    # 绘制结果
    plot_results(train_loss, train_acc, test_acc, epoch_num)
    
    # 显示样本预测
    show_predictions(model, test_dataset)
    
    print(f"\n训练完成！最佳测试准确率: {max(test_acc):.2f}%")