DAY 43 复习日

@浙大疏锦行作业：

kaggle找到一个图像数据集，用cnn网络进行训练并且用grad-cam做可视化

进阶：并拆分成多个文件

我找的是kaggle上的猫狗二分类训练集

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import os
import warnings
from torch.utils.data import random_split
warnings.filterwarnings("ignore")

# 设置中文字体支持
plt.rcParams["font.family"] = ["SimHei", "DejaVu Sans"]
plt.rcParams['axes.unicode_minus'] = False

# 设置随机种子
torch.manual_seed(42)
np.random.seed(42)

# 检查GPU可用性
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 数据预处理和增强
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 加载数据集
train_dataset = datasets.ImageFolder(
    root=r'D:\BaiduNetdiskDownload\kaggle\kaggle\train\train',
    transform=train_transform
)


trainloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=64,
    shuffle=True,
    num_workers=0
)
# 加载数据集
test_dataset = datasets.ImageFolder(
    root=r'D:\BaiduNetdiskDownload\kaggle\kaggle\test1',
    transform=test_transform
)
testloader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=64,
    shuffle=False,  # 测试集不需要打乱
    num_workers=0
)

# 打印数据集信息
print(f"训练集大小: {len(train_dataset)}")
print(f"类别数量: {len(train_dataset.classes)}")
print(f"类别名称: {train_dataset.classes}")
print(f"测试集大小: {len(test_dataset)}")

# 修正后的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self, num_classes=2):  # 猫狗二分类，所以num_classes=2
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)

        # 修正全连接层输入尺寸
        # 输入图像224x224，经过3次池化后变为28x28 (224/2^3=28)
        self.fc1 = nn.Linear(128 * 28 * 28, 512)
        self.fc2 = nn.Linear(512, num_classes)  # 输出类别数改为2
        self.dropout = nn.Dropout(0.5)  # 添加dropout防止过拟合

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))

        # 展平特征图
        x = x.view(-1, 128 * 28 * 28)
        x = self.dropout(F.relu(self.fc1(x)))
        x = self.fc2(x)
        return x


# 初始化模型
model = SimpleCNN(num_classes=len(train_dataset.classes))
model = model.to(device)
print("模型已创建")


# 评估模型函数 - 先定义这个函数
def evaluate_model(model):
    model.eval()  # 设置模型为评估模式
    correct = 0
    total = 0

    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()

    accuracy = 100 * correct / total
    model.train()  # 恢复训练模式
    return accuracy


# 训练模型
def train_model(model, epochs=20):  # 增加训练轮数
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 学习率调度器
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

    model.train()  # 设置模型为训练模式

    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 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()

            # 计算准确率
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            if i % 100 == 99:
                print(f'[{epoch + 1}, {i + 1}] 损失: {running_loss / 100:.3f}, 准确率: {100 * correct / total:.2f}%')
                running_loss = 0.0

        # 更新学习率
        scheduler.step()

        # 每个epoch结束后在测试集上评估
        test_accuracy = evaluate_model(model)
        print(f'Epoch {epoch + 1} 测试集准确率: {test_accuracy:.2f}%')

    print("训练完成")
    # 保存模型
    torch.save(model.state_dict(), 'catdog_cnn.pth')
    print("模型已保存为 'catdog_cnn.pth'")


# 训练模型
print("开始训练模型...")
train_model(model, epochs=20)

# 加载最佳模型
model.load_state_dict(torch.load('catdog_cnn.pth'))
model.eval()
print("已加载训练好的模型")


# Grad-CAM实现
class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        self.register_hooks()

    def register_hooks(self):
        def forward_hook(module, input, output):
            self.activations = output.detach()

        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0].detach()

        self.target_layer.register_forward_hook(forward_hook)
        self.target_layer.register_backward_hook(backward_hook)

    def generate_cam(self, input_image, target_class=None):
        model_output = self.model(input_image)

        if target_class is None:
            target_class = torch.argmax(model_output, dim=1).item()

        self.model.zero_grad()
        one_hot = torch.zeros_like(model_output)
        one_hot[0, target_class] = 1
        model_output.backward(gradient=one_hot)

        gradients = self.gradients
        activations = self.activations

        weights = torch.mean(gradients, dim=(2, 3), keepdim=True)
        cam = torch.sum(weights * activations, dim=1, keepdim=True)
        cam = F.relu(cam)

        # 调整到输入图像大小 (224x224)
        cam = F.interpolate(cam, size=(224, 224), mode='bilinear', align_corners=False)
        cam = cam - cam.min()
        cam = cam / cam.max() if cam.max() > 0 else cam

        return cam.cpu().squeeze().numpy(), target_class


# 选择目标层进行Grad-CAM可视化
target_layer = model.conv3  # 使用最后一个卷积层

# 创建Grad-CAM实例
grad_cam = GradCAM(model, target_layer)

# 获取一个测试样本进行可视化
dataiter = iter(testloader)
images, labels = next(dataiter)

# 选择第一个图像进行可视化
image = images[0:1].to(device)  # 保持batch维度
true_label = labels[0].item()

# 生成热图
heatmap, pred_class = grad_cam.generate_cam(image)


# 可视化结果
def visualize_gradcam(image, heatmap, true_label, pred_label, class_names):
    # 将图像转换为numpy数组
    image = image.cpu().squeeze().numpy()
    image = np.transpose(image, (1, 2, 0))

    # 反归一化
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = std * image + mean
    image = np.clip(image, 0, 1)

    # 创建子图
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5))

    # 显示原始图像
    ax1.imshow(image)
    ax1.set_title(f'原始图像\n真实类别: {class_names[true_label]}')
    ax1.axis('off')

    # 显示热图
    ax2.imshow(heatmap, cmap='jet')
    ax2.set_title('Grad-CAM热图')
    ax2.axis('off')

    # 显示叠加效果
    heatmap_resized = np.uint8(255 * heatmap)
    heatmap_colored = plt.cm.jet(heatmap_resized)[..., :3]  # 只取RGB，忽略alpha
    superimposed = 0.5 * image + 0.5 * heatmap_colored
    ax3.imshow(superimposed)
    ax3.set_title(f'叠加效果\n预测类别: {class_names[pred_class]}')
    ax3.axis('off')

    plt.tight_layout()
    plt.show()


# 获取类别名称
class_names = train_dataset.classes

# 可视化Grad-CAM结果
visualize_gradcam(images[0], heatmap, true_label, pred_class, class_names)

我在自己的电脑上跑的，速度很慢，所以我只跑了20个epoch，接下来我会增加epoch继续对模型进行训练，看效果会不会好一点