作业：

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

进阶：并拆分成多个文件

从谷歌图片中拍摄的 10 种不同类别的动物图片

数据预处理

import os
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split

def load_data(data_dir, batch_size):
    # 数据预处理
    data_transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # 加载数据集
    image_dataset = datasets.ImageFolder(data_dir, data_transform)

    # 划分训练集和验证集
    train_size = int(0.8 * len(image_dataset))
    val_size = len(image_dataset) - train_size
    train_dataset, val_dataset = random_split(image_dataset, [train_size, val_size])

    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    dataloaders = {'train': train_dataloader, 'val': val_dataloader}
    dataset_sizes = {'train': train_size, 'val': val_size}
    class_names = image_dataset.classes

    return dataloaders, dataset_sizes, class_names

构建并训练 CNN 模型

import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super(SimpleCNN, self).__init__()
        # 定义特征提取层
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        # 定义分类层
        self.classifier = nn.Sequential(
            nn.Linear(64 * 28 * 28, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        # 前向传播，先通过特征提取层，再通过分类层
        x = self.features(x)
        x = x.view(-1, 64 * 28 * 28)
        x = self.classifier(x)
        return x

模型训练模块

import torch
import torch.nn as nn
import torch.optim as optim

def train_model(model, dataloaders, dataset_sizes, criterion, optimizer, num_epochs=25):
    # 判断是否有可用的 GPU，若有则使用 GPU 进行训练
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)

    for epoch in range(num_epochs):
        print(f'第 {epoch} 个 epoch，共 {num_epochs - 1} 个 epochs')
        print('-' * 10)

        # 每个 epoch 都有一个训练和验证阶段
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # 训练模式
            else:
                model.eval()   # 评估模式

            running_loss = 0.0
            running_corrects = 0

            # 迭代数据
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 零参数梯度
                optimizer.zero_grad()

                # 前向传播
                # 只有在训练时才跟踪历史
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # 只有在训练阶段才进行反向传播和优化
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # 统计
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f'{phase} 阶段：损失值: {epoch_loss:.4f} 准确率: {epoch_acc:.4f}')

    return model

Grad-CAM可视化模块 

import torch
import torch.nn.functional as F
import numpy as np
import cv2

class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None

        # 反向传播钩子函数，用于捕获梯度
        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0]

        # 前向传播钩子函数，用于捕获激活值
        def forward_hook(module, input, output):
            self.activations = output

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

    def forward(self, input_tensor):
        # 将模型设置为评估模式并进行前向传播
        self.model.eval()
        output = self.model(input_tensor)
        return output

    def generate_cam(self, input_tensor, target_class=None):
        # 进行前向传播
        output = self.forward(input_tensor)

        # 如果未指定目标类别，则选择输出概率最大的类别
        if target_class is None:
            target_class = torch.argmax(output, dim=1).item()

        one_hot = torch.zeros_like(output)
        one_hot[:, target_class] = 1
        one_hot.requires_grad_(True)

        # 清零模型参数的梯度
        self.model.zero_grad()
        # 计算损失并进行反向传播
        (one_hot * output).sum().backward(retain_graph=True)

        gradients = self.gradients[0]
        activations = self.activations[0]

        # 对梯度进行全局平均池化
        pooled_gradients = torch.mean(gradients, dim=[1, 2])

        for i in range(activations.shape[0]):
            activations[i, :, :] *= pooled_gradients[i]

        # 对激活值求和生成 CAM 图
        cam = torch.sum(activations, dim=0).detach().cpu().numpy()
        # 取 CAM 图的正值部分
        cam = np.maximum(cam, 0)
        # 调整 CAM 图的大小以匹配输入图像
        cam = cv2.resize(cam, (input_tensor.shape[3], input_tensor.shape[2]))
        # 归一化 CAM 图
        cam = cam - np.min(cam)
        cam = cam / np.max(cam)
        return cam

 主程序

from data_loader import load_data
from model import SimpleCNN
from train import train_model
from grad_cam import GradCAM
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from PIL import Image
import numpy as np
import cv2
import sys
import os
# 将当前目录添加到 Python 模块搜索路径中
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

if __name__ == '__main__':
    # 加载数据，设置批次大小，你可以根据需要调整该值
    batch_size = 32
    # 修改解包操作以处理所有返回值
    dataloaders, dataset_sizes, class_names = load_data('raw-img', batch_size)

    # 获取类别数量
    num_classes = len(class_names)

    # 使用类别数量初始化模型
    model = SimpleCNN(num_classes)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 训练模型
    trained_model = train_model(model, dataloaders, dataset_sizes, criterion, optimizer, num_epochs=5)

    # 生成 Grad - CAM 可视化结果
    # 修改此处，选择实际存在的卷积层
    # grad_cam = GradCAM(model, target_layer=model.conv2)
    grad_cam = GradCAM(model, target_layer=model.features[0])
    img_path = 'path/to/your/image.jpg'
    img = Image.open(img_path).convert('RGB')
    cam = grad_cam(img)
    plt.imshow(img)
    plt.imshow(cam, alpha=0.5, cmap='jet')
    plt.axis('off')
    plt.savefig('grad_cam_result.jpg')
    plt.show()

 @浙大疏锦行