mobileNetV2.py

import torch
import torch.nn as nn #神经网络核心模块
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import os
import time
import matplotlib.pyplot as plt # 用于绘制损失曲线
import numpy as np
from PIL import Image # 导入 PIL 库用于加载单张测试图像


def main():
    # --- 1. 硬件和路径配置 ---
    # 检查是否有可用的 GPU (RTX 3050)，否则使用 CPU
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"当前使用设备: {DEVICE}")

    # 数据集在 'data/' 目录下
    # 文件夹结构 (train/label 1/image, val/label 1/image)
    TRAIN_DIR = 'data/train'
    VAL_DIR = 'data/val'



    # --- 2. 训练超参数配置 ---
    NUM_CLASSES = 20      # ImageNet 20 类子集
    BATCH_SIZE = 16       # RTX 3050 4GB 显存，选16 或 8
    LEARNING_RATE = 0.001 # 初始学习率
    NUM_EPOCHS = 80       #训练轮次



    # --- 3 数据预处理与加载 ---

    # 定义数据预处理
    # 训练集：需要数据增强（随机裁剪、翻转）来防止过拟合
    train_transforms = transforms.Compose([
        transforms.RandomResizedCrop(224),   # 随机裁剪并缩放到 224x224
        transforms.RandomHorizontalFlip(),   # 随机水平翻转
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), # 颜色抖动 saturationn：饱和度
        transforms.ToTensor(),
        # 转换为张量。
        # 像素值从[0, 255] -> [0.0, 1.0]；形状 (Height, Weight, Channel)-> （C,H,W)
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        # ImageNet 标准归一化。是在整个 ImageNet 数据集上统计得到的均值和标准差
    ])

    # 验证集：不需要数据增强，只需居中裁剪和归一化
    val_transforms = transforms.Compose([
        transforms.Resize(256),              # 先缩放到 256
        transforms.CenterCrop(224),          # 再中心裁剪到 224x224
        transforms.ToTensor(),      # 转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 使用 ImageFolder 加载数据集
    # ImageFolder 会自动根据 'data/train' 下的子文件夹名称 (label 1, label 2...) 分配标签
    train_dataset = datasets.ImageFolder(root=TRAIN_DIR, transform=train_transforms)
    val_dataset = datasets.ImageFolder(root=VAL_DIR, transform=val_transforms)

    # 创建 DataLoader
    # num_workers 可以设置多线程加载数据，加快速度
    train_loader = DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
    val_loader = DataLoader(dataset=val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)

    print(f"数据加载完成。训练集样本数: {len(train_dataset)}, 验证集样本数: {len(val_dataset)}")
    print(f"类别: {train_dataset.classes}") # 打印自动检测到的类别



    #  ------- 4、模型的定义 -------
    #  ---- 4.1 倒残差块（InvertedResidual） ----
    class InvertedResidual(nn.Module):
        """
        实现 MobileNetV2 的核心：倒残差块
        结构: 1x1 升维 -> 3x3 深度卷积 -> 1x1 线性降维
        """

        def __init__(self, inp_channels, out_channels, stride, expand_ratio):
            # expand_ratio： 第一步升维的时候， 输入通道数扩大倍数
            super(InvertedResidual, self).__init__()
            self.stride = stride

            # 计算隐藏层的通道数（升维后的通道数）
            hidden_channels = int(inp_channels * expand_ratio)

            # 标志：是否使用残差连接 (Shortcut)
            # 仅当步长为1（让图像尺寸不变）且输入输出通道数相同时，才使用（才能做加法）
            self.use_residual = (self.stride == 1) and (inp_channels == out_channels)

            # 核心卷积层
            layers = []

            # 1. 升维 (Expansion): 1x1 逐点卷积
            # 仅当 expand_ratio 不为 1 时才需要升维 (序列2除外)
            if expand_ratio != 1:
                layers.extend([
                    nn.Conv2d(inp_channels, hidden_channels, kernel_size=1, stride=1, padding=0, bias=False),
                    #bias： 卷积层不添加偏置
                    nn.BatchNorm2d(hidden_channels),  # 对每个通道单独 数据归一化+ 缩放与偏移。
                    # Batch Normalization for 2D data  参数：输入通道数
                    nn.ReLU6(inplace=True)  # 使用 ReLU6 激活函数
                    # inplace=True 表示计算结果会 直接修改输入张量的内存空间
                ])

            # 2. 特征过滤 (Depthwise): 3x3 深度卷积
            layers.extend([
                    #可迭代对象添加到列表layers
                nn.Conv2d(hidden_channels, hidden_channels, kernel_size=3, stride=stride, padding=1,
                          groups=hidden_channels, bias=False),
                # groups 表示将输入通道分组，每组单独进行卷积
                # 当 groups=输入通道数 时，每个输入通道仅与一个卷积核卷积（不跨通道混合），极大减少计算量。是深度卷积的关键
                nn.BatchNorm2d(hidden_channels),
                nn.ReLU6(inplace=True)
            ])

            # 3. 降维 (Projection): 1x1 逐点卷积 (线性瓶颈)
            layers.extend([
                nn.Conv2d(hidden_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(out_channels)
                # 注意：这里故意不加激活函数 (ReLU)，这就是 "线性瓶颈"
            ])

            self.conv = nn.Sequential(*layers) #数据一层一层的传。* 在这里是 Python 的解包运算符

        def forward(self, x): #前向传播逻辑
            if self.use_residual:
                # 如果使用残差连接，则将输入 x 与 卷积块的输出 相加
                return x + self.conv(x)
            else:
                # 否则，直接返回卷积块的输出
                return self.conv(x)

    #  ---- 4.2 组装MobileNetV2 完整网络 ----
    class MobileNetV2(nn.Module):
        """
        组装 MobileNetV2 完整网络
        """

        def __init__(self, num_classes=20):
            super(MobileNetV2, self).__init__()

            # MobileNetV2 的网络结构配置表
            # 每一行代表一个序列 (t, c, n, s)
            # t: 扩张倍数, c: 输出通道数, n: 重复次数, s: 序列首块的步长
            self.block_settings = [
                # t,  c, n, s
                [1, 16, 1, 1],  # 序列 2
                [6, 24, 2, 2],  # 序列 3
                [6, 32, 3, 2],  # 序列 4
                [6, 64, 4, 2],  # 序列 5
                [6, 96, 3, 1],  # 序列 6
                [6, 160, 3, 2],  # 序列 7
                [6, 320, 1, 1]  # 序列 8
            ]

            # --- 1. 初始卷积层 (序列 1) ---
            self.initial_conv = nn.Sequential(
                nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1, bias=False),
                nn.BatchNorm2d(32), # 对每个通道单独 数据归一化+ 缩放与偏移。 参数：输入通道数
                nn.ReLU6(inplace=True)
            )

            # --- 2. 倒残差块 (序列 2-8) ---
            self.features = []
            input_channels = 32  # 初始卷积层的输出通道

            # 循环读取配置表，创建倒残差块
            for t, c, n, s in self.block_settings:
                output_channels = c
                # 循环 n 次
                for i in range(n):
                    # 只有序列的第一个块才使用步长 s (用于下采样(缩小特征图尺寸)); 其余块步长为 1，保持特征图尺寸不变
                    stride = s if i == 0 else 1

                    self.features.append(
                        InvertedResidual(input_channels, output_channels, stride, expand_ratio=t)
                    )

                    # 更新下一块的输入通道
                    input_channels = output_channels

            # 将列表 'features' 转换为 nn.Sequential
            self.features = nn.Sequential(*self.features)

            # --- 3. 最后的特征提取层 (序列 9) ---
            self.final_conv = nn.Sequential(
                nn.Conv2d(input_channels, 1280, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(1280),       # # 对每个通道单独 数据归一化+ 缩放与偏移。 参数：输入通道数
                nn.ReLU6(inplace=True)
            )

            # --- 4. 分类头 (序列 10-11) ---
            self.classifier = nn.Sequential(
                nn.AdaptiveAvgPool2d((1, 1)),  # 自适应(只需要指定输出尺寸）平均池化 (序列 10)。 输出形状：(B, 1280, 1, 1)
                nn.Flatten(),  # 展平为 (B, 1280)
                nn.Dropout(0.2),  # Dropout 防止过拟合, 按照0.2的概率随机 “丢弃”（设置为 0）一部分神经元的输出
                nn.Linear(1280, num_classes)  # 全连接层 (序列 11)
                # 用于实现输入特征与输出特征的线性变换：output = input × weight + bias
                #num_classes： 输出的20个类别
            )

        def forward(self, x):
            # 定义数据的前向传播流程
            x = self.initial_conv(x)
            x = self.features(x)
            x = self.final_conv(x)
            x = self.classifier(x)
            return x






    # --- 5. 实例化模型、损失函数和优化器 ---

    # 实例化模型
    model = MobileNetV2(num_classes=NUM_CLASSES)
    model = model.to(DEVICE) # 将模型移动到 GPU

    # 定义损失函数 (CrossEntropyLoss)
    criterion = nn.CrossEntropyLoss()

    # 定义优化器 (AdamW)
    # AdamW 是 Adam 的改进版，通常效果更好
    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE)  #参数1：我们model的所有parameters 参数2：学习率

    # (先不用，试一下) 定义学习率调度器，帮助模型后期更好地收敛
    # scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)

    print("模型和优化器配置完成。")





    # --- 6. 开始训练与验证循环 ---

    # 用于存储损失历史，以便后续绘图
    train_loss_history = []
    val_loss_history = []
    val_accuracy_history = []

    print(f"开始训练，共 {NUM_EPOCHS} 轮...")

    start_time = time.time()

    for epoch in range(NUM_EPOCHS):
        # --- 训练阶段 ---
        #每一轮（Epoch）：
        model.train()  # 设置为训练模式 (BatchNorm, Dropout 生效)

        running_train_loss = 0.0
        for images, labels in train_loader:  #train_loader：前面加载的
            #每一批次图片
            # 1. 将数据移动到 GPU
            images = images.to(DEVICE)
            labels = labels.to(DEVICE)

            # 2. 清零梯度
            optimizer.zero_grad()       #optimizer：前面定义的Adam优化器

            # 3. 前向传播
            outputs = model(images)     #model：前面定义的MobileNetV2模型
            # 等价于调用model.__call__(images)
            # 而 __call__ 内部会调用model.forward(images)

            # 4. 计算损失
            loss = criterion(outputs, labels) #criterion：前面定义的CrossEntropyLoss

            # 5. 反向传播
            loss.backward()

            # 6. 更新权重
            optimizer.step()

            # 累加训练损失
            running_train_loss += loss.item() * images.size(0)
            #loss.item():当前批次的（每个样本）平均损失
            #image:（BatchSize，Channel，H，W）的张量
            # image.size（0）代表第0个维度大小。前面都是16，最后一批不足16

        # (可选) 更新学习率
        # scheduler.step()

        epoch_train_loss = running_train_loss / len(train_dataset)
        train_loss_history.append(epoch_train_loss)

        # --- 验证阶段 ---
        model.eval()  # 设置为评估模式 (BatchNorm, Dropout 冻结)

        running_val_loss = 0.0
        correct_predictions = 0

        with torch.no_grad():  # 在验证阶段，不计算梯度，节省显存
            for images, labels in val_loader:
                #每一批次图片：
                # 1. 将数据移动到 GPU
                images = images.to(DEVICE)
                labels = labels.to(DEVICE)

                # 2. 前向传播
                outputs = model(images)
                '''
                outputs 是模型的输出张量，形状为 (x:BatchSize, y:num_classes类别数（20）)
                eg：outputs = torch.tensor([
                        [1.2, 3.5, 0.8, ...],  # 样本1的20个类别分数
                        [2.1, 0.9, 4.3, ...]   # 样本2的20个类别分数
                        ...                    # 样本n（BatchSize）的20个类别分数
                    ])
                '''

                # 3. 计算损失
                loss = criterion(outputs, labels)
                running_val_loss += loss.item() * images.size(0)

                # 4. 计算准确率
                _, predicted_labels = torch.max(outputs, 1)
                #torch.max(outputs, 1) 返回output 第一个维度（列，类别）之间比较的最大值
                # _占位符，接受最大值。 predicted_labels接受索引

                correct_predictions += (predicted_labels == labels).sum().item()
                #sum是汇总true的个数

        epoch_val_loss = running_val_loss / len(val_dataset)
        val_loss_history.append(epoch_val_loss)

        epoch_val_accuracy = 100.0 * correct_predictions / len(val_dataset)
        val_accuracy_history.append(epoch_val_accuracy)

        # --- 打印 Epoch 结果 ---
        print(f"Epoch [{epoch + 1}/{NUM_EPOCHS}] - "
              f"训练损失: {epoch_train_loss:.4f} - "
              f"验证损失: {epoch_val_loss:.4f} - "
              f"验证准确率: {epoch_val_accuracy:.2f}%")

    end_time = time.time()
    print(f"训练+验证完成。总耗时: {(end_time - start_time) / 60:.2f} 分钟")

    # 保存模型权重
    torch.save(model.state_dict(), 'mobilenetv2_scratch_model.pth')
    # state_dict 是 PyTorch 中用于存储模型可学习参数的字典（dict）
    #路径为相对路径





    # --- 7. 绘制损失函数变化曲线图 ---

    plt.figure(figsize=(10, 5))
    plt.title("模型损失函数变化曲线")
    plt.plot(train_loss_history, label="训练损失 (Train Loss)")
    plt.plot(val_loss_history, label="验证损失 (Validation Loss)")
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.legend()
    plt.grid(True)
    # 保存图像
    plt.savefig("loss_curve.png")
    # 显示图像
    plt.show()




    # --- 8. 可视化卷积核与特征图 ---


    # 验证集中的一张图片路径
    TEST_IMAGE_PATH = 'data/val/n04515003/n0451500300000019.jpg'


    # --- 8.1 可视化卷积核 (Kernels) ---
    # 思路：访问模型第一层 (initial_conv) 的权重，并将其可视化。

    def normalize_kernels(kernels):
        """
        辅助函数：将卷积核权重归一化到 [0, 1] 范围，以便 matplotlib 显示
        """
        # 逐个通道进行归一化，处理维度 (C, H, W)
        min_val = kernels.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
        max_val = kernels.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0]
        # 归一化公式: (x - min) / (max - min)
        normalized_kernels = (kernels - min_val) / (max_val - min_val + 1e-5)
        return normalized_kernels


    # 访问第一层 (initial_conv) 的卷积核（权重）
    # 这一层是 nn.Conv2d(3, 32, ...)，权重形状是 [32, 3, 3, 3]
    kernels = model.initial_conv[0].weight.detach().cpu()

    # 归一化权重以便显示
    kernels = normalize_kernels(kernels)

    # 绘制
    print("\n--- 开始绘制第一层的 32 个卷积核 ---")
    plt.figure(figsize=(16, 8))
    plt.suptitle("MobileNetV2 初始卷积核 (32个)", fontsize=16)
    for i in range(kernels.shape[0]):  # 遍历 32 个核
        plt.subplot(4, 8, i + 1)  # 创建 4x8 的子图网格

        # 将 PyTorch 格式 (C, H, W) 转换为 Matplotlib 格式 (H, W, C)
        kernel_img = kernels[i].permute(1, 2, 0)

        plt.imshow(kernel_img)
        plt.axis('off')
        # plt.title(f'Kernel {i+1}') # 图太多，不显示标题

    plt.savefig("kernels_visualization.png")
    plt.show()

    # --- 8.2 可视化特征图 (Feature Maps) ---
    # 思路：使用 PyTorch 的 Hook 技术捕获模型中间层的输出。

    # 8.2.1 准备测试图像

    if not os.path.exists(TEST_IMAGE_PATH):
        # 仅在默认路径不存在时创建随机张量
        print("警告: 缺少真实测试图片，使用随机张量代替。")
        test_image_tensor = torch.randn(1, 3, 224, 224).to(DEVICE)
    else:
        # 3.3 加载和预处理图像
        try:
            image = Image.open(TEST_IMAGE_PATH).convert('RGB')
            test_image_tensor = val_transforms(image)  # 使用前面定义的 val_transforms
            # 增加一个 Batch 维度 (C, H, W) -> (1, C, H, W)
            test_image_tensor = test_image_tensor.unsqueeze(0).to(DEVICE)
            print(f"已加载测试图片: {TEST_IMAGE_PATH}")
        except Exception as e:
            print(f"错误: 无法加载图片 {TEST_IMAGE_PATH}，请检查路径。错误信息: {e}")
            # 如果加载失败，同样使用随机张量
            test_image_tensor = torch.randn(1, 3, 224, 224).to(DEVICE)

    # 8.2.2 定义 Hook 机制

    feature_maps = []  # 全局列表，用于存储 Hook 捕获的输出


    def hook_fn(module, input, output):
        """
        Hook 函数：捕获目标层的输出 Feature Map 并存储到列表中。
        """
        feature_maps.append(output.detach().cpu())


    # 8.2.3 选择目标层并注册 Hook
    # 选择一个较浅层（例如第一个倒残差块）和较深层进行对比：

    # 浅层示例：第一个倒残差块 (model.features[0])
    target_layer_shallow = model.features[0]
    handle_shallow = target_layer_shallow.register_forward_hook(hook_fn)

    # 深层示例：第十个倒残差块 (model.features[9])
    target_layer_deep = model.features[9]
    handle_deep = target_layer_deep.register_forward_hook(hook_fn)

    # 8.2.4 运行模型
    with torch.no_grad():  # 确保不计算梯度
        model(test_image_tensor)

    # 8.2.5 移除 Hook (防止内存泄漏)
    handle_shallow.remove()
    handle_deep.remove()

    # 8.2.6 绘制特征图

    if feature_maps:
        # feature_maps[0] 是浅层的输出，feature_maps[1] 是深层的输出

        for i, (fmap, target_layer) in enumerate(zip(feature_maps, [target_layer_shallow, target_layer_deep])):
            fmap = fmap.squeeze(0)  # 去掉 Batch 维度 -> [C, H, W]

            # 为了清晰，我们只显示前 16 个通道
            num_to_show = min(16, fmap.shape[0])

            layer_name = f"Layer {i + 1}: {target_layer.__class__.__name__}"
            print(f"正在绘制 {layer_name} 的前 {num_to_show} 个特征图 (Feature Maps)...")

            plt.figure(figsize=(12, 12))
            plt.suptitle(f"特征图可视化 ({layer_name})", fontsize=16)

            for j in range(num_to_show):
                plt.subplot(4, 4, j + 1)  # 4x4 网格

                # Feature Map 通常是灰度图
                channel_map = fmap[j, :, :]

                # 使用 cmap='gray'
                plt.imshow(channel_map, cmap='gray')
                plt.axis('off')
                plt.title(f'Channel {j + 1}')

            plt.savefig(f"feature_maps_visualization_layer{i + 1}.png")
            plt.show()

    else:
        print("模型运行失败，未能捕获到特征图。请检查模型是否正确加载。")

    pass

# --- 10. 程序入口 ---
if __name__ == '__main__':
    main()