python打卡day43

作业：

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

进阶：并拆分成多个文件

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集（需根据实际数据集修改路径）
# 假设数据集结构为 root/train/class_x/xxx.png 和 root/test/class_y/xxx.png
# train_dataset = datasets.ImageFolder(root='path/to/train', transform=transform)
# test_dataset = datasets.ImageFolder(root='path/to/test', transform=transform)
# train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
# test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)

# 定义简单的 CNN 模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(32 * 56 * 56, 128)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)  # 假设是 10 分类任务，需根据实际修改

    def forward(self, x):
        x = self.pool1(self.relu1(self.conv1(x)))
        x = self.pool2(self.relu2(self.conv2(x)))
        x = x.view(-1, 32 * 56 * 56)
        x = self.relu3(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}')
from gradcam.utils import visualize_cam
from gradcam import GradCAM
import cv2
import matplotlib.pyplot as plt

# 选择模型的最后一个卷积层
target_layer = model.conv2

# 创建 GradCAM 对象
gradcam = GradCAM(model, target_layer)

# 选择一张测试图像（需根据实际修改）
# image, _ = test_dataset[0]
# image = image.unsqueeze(0)

# 获取 Grad-CAM 结果
mask, _ = gradcam(image)
heatmap, result = visualize_cam(mask, image)

# 转换为 numpy 数组并调整通道顺序
heatmap = heatmap.permute(1, 2, 0).detach().cpu().numpy()
result = result.permute(1, 2, 0).detach().cpu().numpy()

# 可视化结果
plt.figure(figsize=(12, 4))
plt.subplot(131)
plt.imshow(image.squeeze().permute(1, 2, 0).detach().cpu().numpy())
plt.title('Original Image')
plt.subplot(132)
plt.imshow(heatmap)
plt.title('Grad-CAM Heatmap')
plt.subplot(133)
plt.imshow(result)
plt.title('Overlay Result')
plt.show()