深度学习代码解读——自用

肆——2025-03-05 20:13

代码来自：GitHub - ChuHan89/WSSS-Tissue

借助了一些人工智能

2_generate_PM.py

功能总结

该代码用于 生成弱监督语义分割（WSSS）所需的伪掩码（Pseudo-Masks），是 Stage2 训练的前置步骤。其核心流程为：

  1. 加载 Stage1 训练好的分类模型（支持 CAM 生成）。

  2. 为不同层次的特征图生成伪掩码 （如 b4_5, b5_2, bn7 对应的不同网络层）。

  3. 保存伪掩码图像，使用调色板将类别标签映射为彩色图像。

代码解析

1. 导入依赖库
复制代码 
import os
import torch
import argparse
import importlib
from torch.backends import cudnn
cudnn.enabled = True  # 启用CUDA加速
from tool.infer_fun import create_pseudo_mask  # 自定义函数：生成伪掩码

  • 关键依赖：

    • cudnn.enabled = True：启用 cuDNN 加速，优化 GPU 计算性能。

    • create_pseudo_mask：核心函数（用户需参考其实现），负责生成并保存伪掩码。

2. 主函数与参数解析
复制代码 
if __name__ == '__main__':
    # 定义命令行参数
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", default='checkpoints/stage1_checkpoint_trained_on_bcss.pth', type=str)
    parser.add_argument("--network", default="network.resnet38_cls", type=str)
    parser.add_argument("--dataroot", default="datasets/BCSS-WSSS/", type=str)
    parser.add_argument("--dataset", default="bcss", type=str)
    parser.add_argument("--num_workers", default=8, type=int)
    parser.add_argument("--n_class", default=4, type=int)
    args = parser.parse_args()
    print(args)  # 打印参数列表

  • 参数说明：

    • --weights：Stage1 训练好的模型权重文件路径（默认指向 BCSS 数据集）。

    • --network：网络结构定义文件（如 network.resnet38_cls）。

    • --dataroot：数据集根目录（包含训练/测试数据）。

    • --dataset：数据集标识（bcss 或 luad）。

    • --n_class：类别数量（BCSS 为 4 类，LUAD 可能不同）。

3. 定义调色板（颜色映射）
复制代码 
    if args.dataset == 'luad':
        palette = [0]*15  # 初始化长度为15的列表（每类3个RGB通道）
        palette[0:3] = [205,51,51]    # 类别1：红色
        palette[3:6] = [0,255,0]      # 类别2：绿色
        palette[6:9] = [65,105,225]   # 类别3：蓝色
        palette[9:12] = [255,165,0]   # 类别4：橙色
        palette[12:15] = [255, 255, 255]  # 背景或未标注区域：白色
    elif args.dataset == 'bcss':
        palette = [0]*15
        palette[0:3] = [255, 0, 0]    # 类别1：红色
        palette[3:6] = [0,255,0]      # 类别2：绿色
        palette[6:9] = [0,0,255]      # 类别3：蓝色
        palette[9:12] = [153, 0, 255] # 类别4：紫色
        palette[12:15] = [255, 255, 255]  # 背景：白色

  • 作用：将类别标签映射为 RGB 颜色，用于伪掩码的可视化。

  • 细节：

    • 每个类别占 3 个连续位置（RGB 通道）。

    • palette[12:15] 可能表示背景或未标注区域。

    • 不同数据集使用不同的颜色方案（如 BCSS 用紫色表示第4类）。

4. 创建伪掩码保存路径
复制代码 
    PMpath = os.path.join(args.dataroot, 'train_PM')  # 路径示例：datasets/BCSS-WSSS/train_PM
    if not os.path.exists(PMpath):
        os.mkdir(PMpath)  # 若目录不存在则创建
  • 目的 ：在数据集根目录下创建 train_PM 文件夹，用于保存生成的伪掩码。
5. 加载模型
复制代码 
    model = getattr(importlib.import_module("network.resnet38_cls"), 'Net_CAM')(n_class=args.n_class)
    model.load_state_dict(torch.load(args.weights), strict=False)
    model.eval()  # 设置为评估模式（禁用Dropout等随机操作）
    model.cuda()  # 将模型移至GPU

  • 关键步骤：

    • 动态加载模型 ：从 network.resnet38_cls 模块加载 Net_CAM 类（支持 CAM 生成的变体）。

    • 加载权重 ：使用 Stage1 训练好的模型参数（strict=False 允许部分参数不匹配）。

    • 评估模式：关闭 BatchNorm 和 Dropout 的随机性，确保结果一致性。

6. 生成多级伪掩码
复制代码 
    ##
    fm = 'b4_5'  # 特征模块名称（可能对应网络中的某个中间层）
    savepath = os.path.join(PMpath, 'PM_' + fm)  # 保存路径：train_PM/PM_b4_5
    if not os.path.exists(savepath):
        os.mkdir(savepath)
    create_pseudo_mask(model, args.dataroot, fm, savepath, args.n_class, palette, args.dataset)

    ## 重复相同流程生成其他层级的伪掩码
    fm = 'b5_2'
    savepath = os.path.join(PMpath, 'PM_' + fm)
    if not os.path.exists(savepath):
        os.mkdir(savepath)
    create_pseudo_mask(model, args.dataroot, fm, savepath, args.n_class, palette, args.dataset)

    ##
    fm = 'bn7'
    savepath = os.path.join(PMpath, 'PM_' + fm)
    if not os.path.exists(savepath):
        os.mkdir(savepath)
    create_pseudo_mask(model, args.dataroot, fm, savepath, args.n_class, palette, args.dataset)

  • 功能 ：针对不同特征模块（fm）生成伪掩码，保存到对应子目录。

  • 关键参数：

    • fm：特征模块标识，可能对应网络中的不同层（如 ResNet 的 block4、block5 或 bottleneck）。

    • create_pseudo_mask：核心函数，推测其功能为：

      1. 加载训练集图像。

      2. 使用模型提取指定层的特征图。

      3. 生成类别激活图（CAM）。

      4. 根据阈值将 CAM 转换为二值伪掩码。

      5. 应用调色板将掩码保存为彩色 PNG 图像。

代码执行示例

复制代码 
python generate_pseudo_masks.py \
    --dataset bcss \
    --dataroot datasets/BCSS-WSSS/ \
    --weights checkpoints/stage1_checkpoint_trained_on_bcss.pth

  • 输出 ：在 datasets/BCSS-WSSS/train_PM/ 下生成三个子目录：

    • PM_b4_5：基于 b4_5 层特征的伪掩码。

    • PM_b5_2：基于 b5_2 层特征的伪掩码。

    • PM_bn7：基于 bn7 层特征的伪掩码。

总结

该代码是弱监督语义分割流程中 生成多级伪掩码的关键步骤，利用 Stage1 训练的分类模型提取不同层级的特征，生成伪标签供 Stage2 的分割模型训练。通过多级伪掩码的融合，可以提升最终分割结果的精度和鲁棒性。

3_train_stage2.py

功能总结

该代码是弱监督语义分割（WSSS）的 Stage2 训练与测试脚本，核心功能为：

  1. 训练分割模型：基于 DeepLab v3+ 架构，使用 Stage1 生成的伪掩码（Pseudo-Masks）进行监督训练。

  2. 验证与测试：评估模型在验证集和测试集上的性能（如 mIoU、像素准确率等）。

  3. 门控机制（Gate Mechanism）：在测试阶段结合 Stage1 的分类结果过滤分割预测，提升精度。

  4. 多任务损失：融合不同层次伪掩码的损失（主伪掩码 + 两种增强版本）。

代码结构

复制代码 
# 1. 依赖库导入
import argparse, os, numpy as np
from tqdm import tqdm
import torch
from tool.GenDataset import make_data_loader
from network.sync_batchnorm.replicate import patch_replication_callback
from network.deeplab import *
from tool.loss import SegmentationLosses
from tool.lr_scheduler import LR_Scheduler
from tool.saver import Saver
from tool.summaries import TensorboardSummary
from tool.metrics import Evaluator

# 2. 定义训练器类
class Trainer(object):
    def __init__(self, args): ...  # 初始化模型、数据、优化器等
    def training(self, epoch): ...  # 训练一个epoch
    def validation(self, epoch): ...  # 验证集评估
    def test(self, epoch, Is_GM): ...  # 测试集评估（支持门控机制）
    def load_the_best_checkpoint(self): ...  # 加载最佳模型

# 3. 主函数
def main(): ...  # 解析参数、启动训练

if __name__ == "__main__":
    main()

关键代码解析

1. Trainer 类初始化
复制代码 
class Trainer(object):
    def __init__(self, args):
        self.args = args
        # 初始化日志记录与模型保存工具
        self.saver = Saver(args)  # 保存模型检查点
        self.summary = TensorboardSummary('logs')  # TensorBoard日志
        self.writer = self.summary.create_summary()
        # 数据加载
        kwargs = {'num_workers': args.workers, 'pin_memory': False}
        self.train_loader, self.val_loader, self.test_loader = make_data_loader(args, **kwargs)
        # 模型定义（DeepLab v3+）
        self.nclass = args.n_class
        model = DeepLab(
            num_classes=self.nclass,
            backbone=args.backbone,  # 骨干网络（如ResNet）
            output_stride=args.out_stride,  # 输出步长（控制特征图分辨率）
            sync_bn=args.sync_bn,  # 多GPU同步BatchNorm
            freeze_bn=args.freeze_bn  # 冻结BN层参数
        )
        # 优化器配置（分层学习率）
        train_params = [
            {'params': model.get_1x_lr_params(), 'lr': args.lr},  # 骨干网络低学习率
            {'params': model.get_10x_lr_params(), 'lr': args.lr * 10}  # 分类头高学习率
        ]
        optimizer = torch.optim.SGD(
            train_params, 
            momentum=args.momentum,
            weight_decay=args.weight_decay, 
            nesterov=args.nesterov
        )
        # 损失函数（交叉熵或Focal Loss）
        self.criterion = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode=args.loss_type)
        self.model, self.optimizer = model, optimizer
        # 评估工具（计算mIoU等指标）
        self.evaluator = Evaluator(self.nclass)
        # 学习率调度（Poly策略）
        self.scheduler = LR_Scheduler(
            args.lr_scheduler, 
            args.lr, 
            args.epochs, 
            len(self.train_loader)
        )
        # 加载Stage1的分类模型（用于门控机制）
        model_stage1 = getattr(importlib.import_module('network.resnet38_cls'), 'Net_CAM')(n_class=4)
        resume_stage1 = 'checkpoints/stage1_checkpoint_trained_on_'+str(args.dataset)+'.pth'
        weights_dict = torch.load(resume_stage1)
        model_stage1.load_state_dict(weights_dict)
        self.model_stage1 = model_stage1.cuda()
        self.model_stage1.eval()  # 固定Stage1模型参数
        # GPU并行化
        if args.cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
            patch_replication_callback(self.model)  # 修复多GPU BatchNorm同步问题
            self.model = self.model.cuda()
        # 加载预训练权重（如DeepLab预训练模型）
        if args.resume is not None:
            checkpoint = torch.load(args.resume)
            # 处理分类头权重（微调时保留，否则删除）
            if args.ft:
                self.model.load_state_dict(checkpoint['state_dict'])
                self.optimizer.load_state_dict(checkpoint['optimizer'])
            else:
                del checkpoint['state_dict']['decoder.last_conv.8.weight']
                del checkpoint['state_dict']['decoder.last_conv.8.bias']
                self.model.load_state_dict(checkpoint['state_dict'], strict=False)
        # 初始化最佳mIoU
        self.best_pred = 0.0
2. 训练阶段 training
复制代码 
    def training(self, epoch):
        train_loss = 0.0
        self.model.train()
        tbar = tqdm(self.train_loader)  # 进度条
        num_img_tr = len(self.train_loader)
        for i, sample in enumerate(tbar):
            # 加载数据（图像 + 三个伪掩码）
            image, target, target_a, target_b = sample['image'], sample['label'], sample['label_a'], sample['label_b']
            if self.args.cuda:
                image, target, target_a, target_b = image.cuda(), target.cuda(), target_a.cuda(), target_b.cuda()
            # 调整学习率
            self.scheduler(self.optimizer, i, epoch, self.best_pred)
            self.optimizer.zero_grad()
            # 前向传播
            output = self.model(image)
            # 添加额外通道处理类别4（背景或忽略类）
            one = torch.ones((output.shape[0],1,224,224)).cuda()
            output = torch.cat([output, (100 * one * (target==4).unsqueeze(dim=1)], dim=1)
            # 计算多任务损失（主伪掩码 + 两种增强版本）
            loss_o = self.criterion(output, target)
            loss_a = self.criterion(output, target_a)
            loss_b = self.criterion(output, target_b)
            loss = 0.6*loss_o + 0.2*loss_a + 0.2*loss_b
            # 反向传播
            loss.backward()
            self.optimizer.step()
            # 统计损失
            train_loss += loss.item()
            tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
            # 记录TensorBoard日志
            self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_img_tr * epoch)
        # 输出epoch总结
        self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
        print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0]))
        print('Loss: %.3f' % train_loss)
3. 验证阶段 validation
复制代码 
    def validation(self, epoch):
        self.model.eval()
        self.evaluator.reset()
        tbar = tqdm(self.val_loader, desc='\r')
        test_loss = 0.0
        for i, sample in enumerate(tbar):
            image, target = sample[0]['image'], sample[0]['label']
            if self.args.cuda:
                image, target = image.cuda(), target.cuda()
            with torch.no_grad():
                output = self.model(image)
            # 转换为CPU numpy数组
            pred = output.data.cpu().numpy()
            target = target.cpu().numpy()
            pred = np.argmax(pred, axis=1)
            # 处理类别4（设为忽略类）
            pred[target==4] = 4
            # 更新评估指标
            self.evaluator.add_batch(target, pred)
        # 计算并记录指标
        Acc = self.evaluator.Pixel_Accuracy()
        Acc_class = self.evaluator.Pixel_Accuracy_Class()
        mIoU = self.evaluator.Mean_Intersection_over_Union()
        ious = self.evaluator.Intersection_over_Union()
        FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
        # 输出结果
        print('Validation:')
        print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(Acc, Acc_class, mIoU, FWIoU))
        # 保存最佳模型
        if mIoU > self.best_pred:
            self.best_pred = mIoU
            self.saver.save_checkpoint({
                'state_dict': self.model.module.state_dict(),
                'optimizer': self.optimizer.state_dict()
            }, 'stage2_checkpoint_trained_on_'+self.args.dataset+'.pth')
4. 测试阶段 test（含门控机制）
复制代码 
    def test(self, epoch, Is_GM):
        self.load_the_best_checkpoint()  # 加载最佳模型
        self.model.eval()
        self.evaluator.reset()
        tbar = tqdm(self.test_loader, desc='\r')
        for i, sample in enumerate(tbar):
            image, target = sample[0]['image'], sample[0]['label']
            if self.args.cuda:
                image, target = image.cuda(), target.cuda()
            with torch.no_grad():
                output = self.model(image)
                # 门控机制：利用Stage1的分类结果过滤分割预测
                if Is_GM:
                    _, y_cls = self.model_stage1.forward_cam(image)  # Stage1的分类输出
                    y_cls = y_cls.cpu().data
                    pred_cls = (y_cls > 0.1)  # 类别存在性判断（阈值0.1）
            # 应用门控机制
            pred = output.data.cpu().numpy()
            if Is_GM:
                pred = pred * pred_cls.unsqueeze(dim=2).unsqueeze(dim=3).numpy()
            # 处理类别4
            pred = np.argmax(pred, axis=1)
            pred[target==4] = 4
            self.evaluator.add_batch(target, pred)
        # 计算并输出指标
        Acc = self.evaluator.Pixel_Accuracy()
        Acc_class = self.evaluator.Pixel_Accuracy_Class()
        mIoU = self.evaluator.Mean_Intersection_over_Union()
        print('Test:')
        print("Acc:{}, Acc_class:{}, mIoU:{}".format(Acc, Acc_class, mIoU))
5. 主函数 main
复制代码 
def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="WSSS Stage2")
    # 模型结构参数
    parser.add_argument('--backbone', default='resnet', choices=['resnet', 'xception', 'drn', 'mobilenet'])
    parser.add_argument('--out-stride', type=int, default=16)  # 输出步长（控制特征图下采样率）
    parser.add_argument('--Is_GM', type=bool, default=True)  # 是否启用门控机制
    # 数据集参数
    parser.add_argument('--dataroot', default='datasets/BCSS-WSSS/')
    parser.add_argument('--dataset', default='bcss')
    parser.add_argument('--n_class', type=int, default=4)
    # 训练超参数
    parser.add_argument('--epochs', type=int, default=30)
    parser.add_argument('--batch-size', type=int, default=20)
    parser.add_argument('--lr', type=float, default=0.01)
    parser.add_argument('--lr-scheduler', default='poly', choices=['poly', 'step', 'cos'])
    # 其他配置
    parser.add_argument('--gpu-ids', default='0')  # 指定使用的GPU
    parser.add_argument('--resume', default='init_weights/deeplab-resnet.pth.tar')  # 预训练权重
    args = parser.parse_args()
    
    # 配置CUDA
    args.cuda = not args.no_cuda and torch.cuda.is_available()
    if args.cuda:
        args.gpu_ids = [int(s) for s in args.gpu_ids.split(',')]
    # 自动设置SyncBN
    if args.sync_bn is None:
        args.sync_bn = True if args.cuda and len(args.gpu_ids) > 1 else False
    
    # 初始化训练器并启动训练
    trainer = Trainer(args)
    for epoch in range(trainer.args.epochs):
        trainer.training(epoch)
        if epoch % args.eval_interval == 0:
            trainer.validation(epoch)
    # 最终测试
    trainer.test(epoch, args.Is_GM)
    trainer.writer.close()

关键设计解析

  1. 多任务损失：

    • 目标 ：同时优化主伪掩码（target）及其两种增强版本（target_a, target_b），提升模型对不同噪声伪标签的鲁棒性。

    • 权重分配 ：主损失占60%，增强损失各占20%（0.6*loss_o + 0.2*loss_a + 0.2*loss_b）。

  2. 门控机制（Gate Mechanism）：

    • 作用：在测试阶段，利用 Stage1 的分类结果过滤分割预测，仅保留分类模型认为存在的类别。

    • 实现：若 Stage1 对某类别的预测概率 > 0.1，则保留该类的分割结果，否则置零。

  3. 类别4处理：

    • 背景或忽略类 ：在标签中，类别4可能表示背景或未标注区域，预测时直接继承真实标签的值（pred[target==4] = 4），避免错误优化。

  4. 模型初始化：

    • 预训练权重 ：加载 DeepLab 在 ImageNet 上的预训练权重（init_weights/deeplab-resnet.pth.tar），加速收敛。

    • 分层学习率 ：骨干网络使用较低学习率（args.lr），分类头使用更高学习率（args.lr * 10）。

运行示例

复制代码 
python train_stage2.py \
    --dataset bcss \
    --dataroot datasets/BCSS-WSSS/ \
    --backbone resnet \
    --Is_GM True \
    --batch-size 20 \
    --epochs 30

总结

该代码实现了弱监督语义分割的第二阶段训练，通过多任务损失融合多级伪标签，结合门控机制提升测试精度，最终生成高精度分割模型。训练过程支持多GPU加速、Poly学习率调度及多种评估指标监控，适用于医学图像（如BCSS）或自然场景图像的分割任务。

相关推荐
笑衬人心。4 分钟前
JVM 笔记：类加载、内存管理、垃圾收集与垃圾收集器
java·jvm·笔记
无妄之灾|15 分钟前
HCIP笔记（第一、二章）
笔记·hcip
Easy数模15 分钟前
使用llm进行高级主题建模：通过利用 BERTopic 的表示模型和生成式 AI 深入探讨主题建模
人工智能·llm·nlp
22:30Plane-Moon18 分钟前
SQL 查询语法笔记
数据库·笔记·sql
你可以叫我仔哥呀20 分钟前
Java程序员学从0学AI（七）
java·开发语言·人工智能·ai·spring ai
_Kayo_24 分钟前
VUE2 学习笔记8 v-text/html/cloak/once/pre/自定义
笔记·学习
益莱储中国24 分钟前
是德科技 | AI上车后，这条“高速公路”如何畅通？
人工智能·科技
星星火柴93626 分钟前
开发笔记 | 实现人物立绘的差分效果
笔记·unity·游戏程序·优香
普郎特29 分钟前
大白话帮你彻底理解 aiohttp 的 ClientSession 与 ClientResponse 对象
爬虫·python
小张童鞋。30 分钟前
Datawhale AI夏令营--Task2：理解项目目标、从业务理解到技术实现！
人工智能
热门推荐
01Qwen3-Coder 快速上手教程 | Qwen Code + Claude Code02Coze 开源了，送上保姆级私有化部署方案【建议收藏】03扣子开源本地部署教程 丨Coze智能体小白喂饭级指南04全球最强模型Grok4，国内已可免费使用！（附教程）05KGG转MP3工具|非KGM文件|解密音频06腾讯还是太全面了，限时免费！超全CodeBuddy IDE保姆级教程！（附案例）07干翻 Typora！MilkUp：完全免费的桌面端 Markdown 编辑器！0801-开源版COZE-字节 Coze Studio 重磅开源！保姆级本地安装教程，手把手带你体验09vue数据变化但页面不变10ChatGPT Agent 完全使用指南：2025年7月最新功能详解