import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
import matplotlib.pyplot as plt
import numpy as np# 设置中文字体支持，避免绘图时中文乱码
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams['axes.unicode_minus'] = False  # 检查 GPU 是否可用，优先使用 GPU 加速训练
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")# 1. 数据预处理（调整为 128×128 尺寸，适配数据增强）
# 训练集数据增强：先随机裁剪（从原图多样本中截取），再 resize 到 128×128；或直接随机裁剪到目标尺寸（若原图足够大）
# 这里演示更灵活的方式：先随机从原图取一块，再统一成 128×128（假设原图有足够尺寸支撑裁剪，否则需调整策略）
train_transform = transforms.Compose([transforms.RandomResizedCrop(128, scale=(0.6, 1.0)),  # 随机裁剪并 resize 到 128×128，scale 控制裁剪区域占原图比例transforms.RandomHorizontalFlip(),  # 随机水平翻转transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 颜色抖动transforms.RandomRotation(15),  # 随机旋转transforms.ToTensor(),  # 转为张量transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))  # 标准化（可根据实际数据统计均值、方差优化）
])# 测试集预处理：直接 resize 到 128×128，保证尺寸统一
test_transform = transforms.Compose([transforms.Resize((128, 128)),  # 统一 resize 到 128×128transforms.ToTensor(),transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])# 2、用全部数据拆分（示例按 8:2 拆分训练集和测试集）
full_dataset = datasets.ImageFolder(root=r"D:\python_learning\cyl_python\day43CNN_kaggle\BengaliFishImages\fish_images",  transform=train_transform
)
train_size = int(0.8 * len(full_dataset))
test_size = len(full_dataset) - train_size
train_dataset, test_dataset = random_split(full_dataset, [train_size, test_size])# 3. 创建数据加载器，按批次加载数据
batch_size = 32  # 因图像尺寸变大，可适当减小 batch_size 避免显存不足
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)# 4. 定义适配 128×128 尺寸的 CNN 模型（需重新计算全连接层输入维度）
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()  # 第一个卷积块：输入 3 通道，输出 32 通道self.conv1 = nn.Conv2d(3, 32, 3, padding=1)self.bn1 = nn.BatchNorm2d(32)self.relu1 = nn.ReLU()self.pool1 = nn.MaxPool2d(2, 2)  # 128 -> 64# 第二个卷积块：输入 32 通道，输出 64 通道self.conv2 = nn.Conv2d(32, 64, 3, padding=1)self.bn2 = nn.BatchNorm2d(64)self.relu2 = nn.ReLU()self.pool2 = nn.MaxPool2d(2)  # 64 -> 32# 第三个卷积块：输入 64 通道，输出 128 通道self.conv3 = nn.Conv2d(64, 128, 3, padding=1)self.bn3 = nn.BatchNorm2d(128)self.relu3 = nn.ReLU()self.pool3 = nn.MaxPool2d(2)  # 32 -> 16# 第四个卷积块（可选，进一步提取特征，若觉得网络浅可添加）self.conv4 = nn.Conv2d(128, 256, 3, padding=1)self.bn4 = nn.BatchNorm2d(256)self.relu4 = nn.ReLU()self.pool4 = nn.MaxPool2d(2)  # 16 -> 8# 计算全连接层输入维度：根据卷积后输出尺寸推算# 经过上述 4 次池化（若用 3 次，需对应调整），假设用 4 次池化，最终特征图尺寸是 8×8# 若卷积块数量变化，需重新计算：比如 3 次池化，尺寸是 128/(2^3)=16，那维度是 128*16*16 self.fc1 = nn.Linear(256 * 8 * 8, 512)  # 这里按 4 次池化到 8×8 计算self.dropout = nn.Dropout(p=0.5)self.fc2 = nn.Linear(512, 20)  # 20 分类，匹配数据集类别数def forward(self, x):# 卷积块 1 前向x = self.conv1(x)x = self.bn1(x)x = self.relu1(x)x = self.pool1(x)# 卷积块 2 前向x = self.conv2(x)x = self.bn2(x)x = self.relu2(x)x = self.pool2(x)# 卷积块 3 前向x = self.conv3(x)x = self.bn3(x)x = self.relu3(x)x = self.pool3(x)# 卷积块 4 前向（若添加了此块）x = self.conv4(x)x = self.bn4(x)x = self.relu4(x)x = self.pool4(x)# 展平特征图，进入全连接层x = x.view(-1, 256 * 8 * 8)  # 与 __init__ 中 fc1 输入维度对应x = self.fc1(x)x = self.relu3(x)x = self.dropout(x)x = self.fc2(x)return x# 初始化模型，并移动到 GPU（若可用）
model = CNN().to(device)# 定义损失函数、优化器、学习率调度器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=3, factor=0.5
)# 5. 训练与测试函数（复用逻辑，无需修改）
def train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs):model.train()  # 设置为训练模式all_iter_losses = []  # 记录每个 batch 的损失iter_indices = []     # 记录 iteration 序号train_acc_history = []  # 训练集准确率历史test_acc_history = []   # 测试集准确率历史train_loss_history = [] # 训练集损失历史test_loss_history = []  # 测试集损失历史for epoch in range(epochs):running_loss = 0.0correct = 0total = 0for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device)  # 数据移到 GPUoptimizer.zero_grad()  # 梯度清零output = model(data)   # 前向传播loss = criterion(output, target)  # 计算损失loss.backward()        # 反向传播optimizer.step()       # 更新参数# 记录当前 iteration 损失iter_loss = loss.item()all_iter_losses.append(iter_loss)iter_indices.append(epoch * len(train_loader) + batch_idx + 1)# 统计训练集准确率running_loss += iter_loss_, predicted = output.max(1)total += target.size(0)correct += predicted.eq(target).sum().item()# 每 100 个批次打印训练信息if (batch_idx + 1) % 100 == 0:print(f'Epoch: {epoch+1}/{epochs} | Batch: {batch_idx+1}/{len(train_loader)} 'f'| 单Batch损失: {iter_loss:.4f} | 累计平均损失: {running_loss/(batch_idx+1):.4f}')# 计算当前 epoch 训练集指标epoch_train_loss = running_loss / len(train_loader)epoch_train_acc = 100. * correct / totaltrain_acc_history.append(epoch_train_acc)train_loss_history.append(epoch_train_loss)# 测试阶段model.eval()  # 切换为评估模式test_loss = 0correct_test = 0total_test = 0with torch.no_grad():  # 测试时无需计算梯度for data, target in test_loader:data, target = data.to(device), target.to(device)output = model(data)test_loss += criterion(output, target).item()_, predicted = output.max(1)total_test += target.size(0)correct_test += predicted.eq(target).sum().item()# 计算当前 epoch 测试集指标epoch_test_loss = test_loss / len(test_loader)epoch_test_acc = 100. * correct_test / total_testtest_acc_history.append(epoch_test_acc)test_loss_history.append(epoch_test_loss)# 更新学习率调度器scheduler.step(epoch_test_loss)print(f'Epoch {epoch+1}/{epochs} 完成 | 训练准确率: {epoch_train_acc:.2f}% | 测试准确率: {epoch_test_acc:.2f}%')# 绘制 iteration 损失曲线plot_iter_losses(all_iter_losses, iter_indices)# 绘制 epoch 指标曲线plot_epoch_metrics(train_acc_history, test_acc_history, train_loss_history, test_loss_history)return epoch_test_acc# 6. 绘图函数（可视化训练过程，无需修改）
def plot_iter_losses(losses, indices):plt.figure(figsize=(10, 4))plt.plot(indices, losses, 'b-', alpha=0.7, label='Iteration Loss')plt.xlabel('Iteration（Batch序号）')plt.ylabel('损失值')plt.title('每个 Iteration 的训练损失')plt.legend()plt.grid(True)plt.tight_layout()plt.show()def plot_epoch_metrics(train_acc, test_acc, train_loss, test_loss):epochs = range(1, len(train_acc) + 1)plt.figure(figsize=(12, 4))# 绘制准确率曲线plt.subplot(1, 2, 1)plt.plot(epochs, train_acc, 'b-', label='训练准确率')plt.plot(epochs, test_acc, 'r-', label='测试准确率')plt.xlabel('Epoch')plt.ylabel('准确率 (%)')plt.title('训练和测试准确率')plt.legend()plt.grid(True)# 绘制损失曲线plt.subplot(1, 2, 2)plt.plot(epochs, train_loss, 'b-', label='训练损失')plt.plot(epochs, test_loss, 'r-', label='测试损失')plt.xlabel('Epoch')plt.ylabel('损失值')plt.title('训练和测试损失')plt.legend()plt.grid(True)plt.tight_layout()plt.show()# 7. 执行训练（根据显存情况调整 epochs 和 batch_size ）
epochs = 80  # 图像尺寸大，训练慢，可先小批次试训
print("开始使用 CNN 训练自定义鱼类数据集（128×128 尺寸）...")
final_accuracy = train(model, train_loader, test_loader, criterion, optimizer, scheduler, device, epochs)
print(f"训练完成！最终测试准确率: {final_accuracy:.2f}%")# 保存模型（可选，根据需求决定是否保存）
# torch.save(model.state_dict(), 'bengali_fish_cnn_128.pth')

import warnings
warnings.filterwarnings("ignore")
import matplotlib.pyplot as plt
import numpy as np
import torch
from torchvision import transforms
from PIL import Image# 设置中文字体支持
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题# Grad-CAM实现
class GradCAM:def __init__(self, model, target_layer):self.model = modelself.target_layer = target_layerself.feature_maps = Noneself.gradient = None# 注册钩子self.hook_handles = []# 保存特征图的正向钩子def forward_hook(module, input, output):self.feature_maps = output.detach()# 保存梯度的反向钩子def backward_hook(module, grad_in, grad_out):self.gradient = grad_out[0].detach()self.hook_handles.append(target_layer.register_forward_hook(forward_hook))self.hook_handles.append(target_layer.register_backward_hook(backward_hook))# 设置为评估模式self.model.eval()def __call__(self, x):# 前向传播x = x.to(device)self.model.zero_grad()output = self.model(x)# 获取预测类别pred_class = torch.argmax(output, dim=1).item()# 反向传播one_hot = torch.zeros_like(output)one_hot[0, pred_class] = 1output.backward(gradient=one_hot, retain_graph=True)# 计算权重 (全局平均池化梯度)weights = torch.mean(self.gradient, dim=(2, 3), keepdim=True)# 加权组合特征图cam = torch.sum(weights * self.feature_maps, dim=1).squeeze()# ReLU激活，因为我们只关心对预测有积极贡献的区域cam = torch.relu(cam)# 归一化if torch.max(cam) > 0:cam = cam / torch.max(cam)# 调整为图像大小cam = torch.nn.functional.interpolate(cam.unsqueeze(0).unsqueeze(0),size=(128, 128),  # 匹配图像尺寸mode='bilinear',align_corners=False).squeeze()return cam.cpu().numpy(), pred_classdef remove_hooks(self):for handle in self.hook_handles:handle.remove()# 转换图像以便可视化
def tensor_to_np(tensor):img = tensor.cpu().numpy().transpose(1, 2, 0)# 使用实际训练时的均值和标准差mean = np.array([0.4914, 0.4822, 0.4465])std = np.array([0.2023, 0.1994, 0.2010])img = std * img + mean  # 反标准化img = np.clip(img, 0, 1)  # 确保像素值在[0,1]范围内return img# 选择一个随机图像
# idx = np.random.randint(len(test_dataset))
idx = 110  # 选择测试集中的第103张图片 (索引从0开始)
image, label = test_dataset[idx]# 获取类别名称（根据ImageFolder自动映射的类别顺序）
classes = test_dataset.dataset.classes  # 注意：如果使用了random_split，需要从原始dataset获取classes
print(f"选择的图像类别: {classes[label]}")# 添加批次维度并移动到设备
input_tensor = image.unsqueeze(0).to(device)# 初始化Grad-CAM（使用model.conv3作为目标层）
# grad_cam = GradCAM(model, model.conv3)
grad_cam = GradCAM(model, model.conv4)# 生成热力图（修正方法调用）
heatmap, pred_class = grad_cam(input_tensor)  # 直接调用对象，而非使用generate_cam方法# 可视化
plt.figure(figsize=(15, 5))# 原始图像
plt.subplot(1, 3, 1)
plt.imshow(tensor_to_np(image))
plt.title(f"原始图像: {classes[label]}")
plt.axis('off')# 热力图
plt.subplot(1, 3, 2)
plt.imshow(heatmap, cmap='jet')
plt.title(f"Grad-CAM热力图: {classes[pred_class]}")
plt.axis('off')# 叠加的图像
plt.subplot(1, 3, 3)
img = tensor_to_np(image)
heatmap_resized = np.uint8(255 * heatmap)
heatmap_colored = plt.cm.jet(heatmap_resized)[:, :, :3]
superimposed_img = heatmap_colored * 0.4 + img * 0.6
plt.imshow(superimposed_img)
plt.title("叠加热力图")
plt.axis('off')plt.tight_layout()
plt.savefig('grad_cam_result_128_crop_100.png')
plt.show()print("Grad-CAM可视化完成。已保存为grad_cam_result_128_crop_100.png")

@浙大疏锦行