图像分割实战-系列教程17：deeplabV3+ VOC分割实战5-------main.py

机器学习杨卓越

已于 2024-01-21 23:16:51 修改

阅读量651

点赞数 11

分类专栏：图像分割实战文章标签：深度学习 pytorch 计算机视觉图像分割 deeplab

于 2024-01-20 17:37:44 首次发布

本文链接：https://blog.csdn.net/weixin_50592077/article/details/135718999

版权

图像分割实战专栏收录该内容

19 篇文章 4 订阅

订阅专栏

在这里插入图片描述

🍁🍁🍁图像分割实战-系列教程总目录

有任何问题欢迎在下面留言
本篇文章的代码运行界面均在Pycharm中进行
本篇文章配套的代码资源已经上传

deeplab系列算法概述
 deeplabV3+ VOC分割实战1
deeplabV3+ VOC分割实战2
deeplabV3+ VOC分割实战3
deeplabV3+ VOC分割实战4
deeplabV3+ VOC分割实战5

10、main.py的main()函数

def main():
    opts = get_argparser().parse_args()
    if opts.dataset.lower() == 'voc':
        opts.num_classes = 21
    elif opts.dataset.lower() == 'cityscapes':
        opts.num_classes = 19

定义main函数
调用参数函数，解析命令行参数
判断数据集名称是否为voc，是则设置类别数为21，21=20个对象+1背景

    # Setup visualization
    vis = Visualizer(port=opts.vis_port, env=opts.vis_env) if opts.enable_vis else None
    if vis is not None:  # display options
        vis.vis_table("Options", vars(opts))
    os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpu_id
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("Device: %s" % device)

设置可视化工具和配置训练设备

根据 opts.enable_vis 的值决定是否创建一个 Visualizer 对象，其端口和环境设置由 opts.vis_port 和 opts.vis_env 提供。Visualizer 是一个python可视化工具类，用于在训练过程中显示图像、图表等信息
检查是否启用了可视化
如果可视化被启用，这一行调用 vis 对象的 vis_table 方法来显示配置选项。vars(opts) 是将 opts 对象转换为字典，其中包含了所有的命令行参数
设置环境变量 CUDA_VISIBLE_DEVICES，其值为 opts.gpu_id
如果GPU可用，则使用 CUDA；否则，使用 CPU
打印出使用的设备信息

    torch.manual_seed(opts.random_seed)
    np.random.seed(opts.random_seed)
    random.seed(opts.random_seed)
    if opts.dataset=='voc' and not opts.crop_val:
        opts.val_batch_size = 1
    train_dst, val_dst = get_dataset(opts)
    train_loader = data.DataLoader( train_dst, batch_size=opts.batch_size, shuffle=True, num_workers=0)
    val_loader = data.DataLoader( val_dst, batch_size=opts.val_batch_size, shuffle=True, num_workers=0)
    print("Dataset: %s, Train set: %d, Val set: %d" % (opts.dataset, len(train_dst), len(val_dst)))

分别为pytorch
numpy
python设置全局随机种子
检查是否使用VOC数据集并且没有启用验证集的裁剪
如果条件满足，将验证批次大小设置为 1
调用 get_dataset 函数来获取训练和验证数据集，该函数在第3部分已经解析
训练集DataLoader，opts.batch_size训练批次大小
验证集DataLoader，opts.val_batch_size验证批次大小
打印出使用的数据集名称以及训练和验证集的大小

    model_map = {
        'deeplabv3_resnet50': network.deeplabv3_resnet50,
        'deeplabv3plus_resnet50': network.deeplabv3plus_resnet50,
        'deeplabv3_resnet101': network.deeplabv3_resnet101,
        'deeplabv3plus_resnet101': network.deeplabv3plus_resnet101,
        'deeplabv3_mobilenet': network.deeplabv3_mobilenet,
        'deeplabv3plus_mobilenet': network.deeplabv3plus_mobilenet
    }
    model = model_map[opts.model](num_classes=opts.num_classes, output_stride=opts.output_stride)
    if opts.separable_conv and 'plus' in opts.model:
        network.convert_to_separable_conv(model.classifier)
    utils.set_bn_momentum(model.backbone, momentum=0.01)

这部分设置网络的参数

定义一个模型映射字典，字典包括本项目可选择的多个网络：
deeplabv3的resnet50
deeplabv3+的resnet50
deeplabv3的resnet101
deeplabv3+的resnet101
deeplabv3的mobilenet
deeplabv3+的mobilenet，这些网络在Network文件夹中使用一定方法构建，在这里直接导入
从预设要选择的网络名称、类别数、输出通道数加载网络
检查是否启用可分离卷积、模型名称包含 ‘plus’
如果条件满足，则对模型的分类器部分应用可分离卷积的转换

调用set_bn_momentum函数，设置批量归一化的动量，set_bn_momentum函数：

def set_bn_momentum(model, momentum=0.1):
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.momentum = momentum

    metrics = StreamSegMetrics(opts.num_classes)
    optimizer = torch.optim.SGD(params=[
        {'params': model.backbone.parameters(), 'lr': 0.1*opts.lr},
        {'params': model.classifier.parameters(), 'lr': opts.lr},
    ], lr=opts.lr, momentum=0.9, weight_decay=opts.weight_decay)
    if opts.lr_policy=='poly':
        scheduler = utils.PolyLR(optimizer, opts.total_itrs, power=0.9)
    elif opts.lr_policy=='step':
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=opts.step_size, gamma=0.1)
    if opts.loss_type == 'focal_loss':
        criterion = utils.FocalLoss(ignore_index=255, size_average=True)
    elif opts.loss_type == 'cross_entropy':
        criterion = nn.CrossEntropyLoss(ignore_index=255, reduction='mean')

评价指标、优化器、学习率、损失函数

StreamSegMetrics类是一个评价指标，用于在语义分割任务中跟踪模型的性能，计算的指标包括总体准确率、平均准确率、频率加权平均准确率、平均交并比（IoU）和每个类别的 IoU，通过计算混淆矩阵和从中派生出的多个指标。实例化出一个对象
设置SGD为优化器，并设置一些参数
设置backbone的学习率，预训练的backbone部分通常需要较小的学习率来微调
设置分类器的学习率
设置优化器的初始全局学习率，设置动量帮助加速SGD并抑制震荡，设置权重衰减用于防止过拟合
如果学习率策略设为 “poly”
则使用多项式衰减
如果学习率策略设为 “step”
则使用阶梯式衰减
根据 opts.loss_type 的值选择合适的损失函数，如果损失函数类型为 “focal_loss”
则使用焦点损失（Focal Loss），这对于处理类别不平衡问题很有效、
如果损失函数类型为 “cross_entropy”
则使用交叉熵损失

	utils.mkdir('checkpoints')
    best_score = 0.0
    cur_itrs = 0
    cur_epochs = 0
    if opts.ckpt is not None and os.path.isfile(opts.ckpt):
        checkpoint = torch.load(opts.ckpt, map_location=torch.device('cpu'))
        model.load_state_dict(checkpoint["model_state"])
        model = nn.DataParallel(model)
        model.to(device)
        if opts.continue_training:
            optimizer.load_state_dict(checkpoint["optimizer_state"])
            scheduler.load_state_dict(checkpoint["scheduler_state"])
            cur_itrs = checkpoint["cur_itrs"]
            best_score = checkpoint['best_score']
            print("Training state restored from %s" % opts.ckpt)
        print("Model restored from %s" % opts.ckpt)
        del checkpoint  # free memory
    else:
        print("[!] Retrain")
        model = nn.DataParallel(model)
        model.to(device)

checkpoints，检查点

创建 “checkpoints” 文件夹，存储训练过程中的模型检查点
初始化最佳分数为0，记录验证集上的最高分数
初始化当前迭代次数为0
初始化当前轮数为0
检查是否提供了检查点文件，并且该文件确实存在
加载检查点文件，意味着无论检查点是在CPU还是GPU上保存的，都会先被加载到CPU内存中
从检查点中恢复模型的状态
使用 DataParallel 来利用多个GPU（如果可用）
模型放入GPU
检查是否需要从检查点继续训练
从检查点恢复优化器的状态
从检查点恢复学习率调度器的状态
从检查点恢复当前迭代次数
从检查点恢复当前最佳分数
打印模型恢复信息
删除检查点变量以释放内存
如果没有提供有效的检查点文件
打印重新训练的信息
使用 DataParallel 来利用多个GPU（如果可用）
模型放入GPU

    vis_sample_id = np.random.randint(0, len(val_loader), opts.vis_num_samples, np.int32) if opts.enable_vis else None
    denorm = utils.Denormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # denormalization for ori images
    if opts.test_only:
        model.eval()
        val_score, ret_samples = validate(
            opts=opts, model=model, loader=val_loader, device=device, metrics=metrics, ret_samples_ids=vis_sample_id)
        print(metrics.to_str(val_score))
        return

训练中的可视化设置、测试模式处理、初始化变量

选择可视化样本：如果启用了可视化则从验证数据加载器中随机选择一定数量的样本进行可视化，否则，vis_sample_id 设为 None
反标准化操作：创建一个反标准化对象，将用于将预处理（标准化）后的图像恢复到原始图像的颜色空间
检查是否仅进行测试，如果是：
将模型设置为评估模式
调用 validate 函数，返回评分、选定的样本
打印评估结果
结束函数执行。

    interval_loss = 0
    while True: #cur_itrs < opts.total_itrs:
        model.train()
        cur_epochs += 1
        for (images, labels) in train_loader:
            cur_itrs += 1
            images = images.to(device, dtype=torch.float32)
            labels = labels.to(device, dtype=torch.long)
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            np_loss = loss.detach().cpu().numpy()
            interval_loss += np_loss

初始化间隔损失为0
开始无限循环（达到最大迭代次数停止）
开启模型训练模式
当前epoch+1
从训练集dataloader取出数据和标签
当前迭代次数+1
训练数据进入GPU
训练标签进入GPU
梯度清零
训练数据进入模型后得到输出
输出和标签通过损失函数计算损失
反向传播，计算损失相对于模型参数的梯度
根据计算的梯度更新模型的权重
将损失从GPU（如果使用）转移到CPU，并转换为NumPy格式
累加损失

            if vis is not None:
                vis.vis_scalar('Loss', cur_itrs, np_loss)
            if (cur_itrs) % 10 == 0:
                interval_loss = interval_loss/10
                print("Epoch %d, Itrs %d/%d, Loss=%f" % (cur_epochs, cur_itrs, opts.total_itrs, interval_loss))
                interval_loss = 0.0

可视化、日志记录部分：

条件可视化：检查是否创建了可视化对象
使用可视化工具来记录当前迭代的损失。 'Loss' 是要记录的数据的名称，cur_itrs 是当前迭代次数，np_loss 是当前迭代的损失值
每当当前迭代次数是10的倍数时执行以下操作：
计算过去10次迭代的平均损失
打印当前的epoch、迭代次数、总迭代次数、平均损失
重置间隔损失，为计算下一个间隔的平均损失做准备

            if (cur_itrs) % opts.val_interval == 0:
                save_ckpt('checkpoints/latest_%s_%s_os%d.pth' % (opts.model, opts.dataset, opts.output_stride))
                print("validation...")
                model.eval()
                val_score, ret_samples = validate(
                    opts=opts, model=model, loader=val_loader, device=device, metrics=metrics, ret_samples_ids=vis_sample_id)
                print(metrics.to_str(val_score))
                if val_score['Mean IoU'] > best_score:  # save best model
                    best_score = val_score['Mean IoU']
                    save_ckpt('checkpoints/best_%s_%s_os%d.pth' % (opts.model, opts.dataset,opts.output_stride))
                if vis is not None:  # visualize validation score and samples
                    vis.vis_scalar("[Val] Overall Acc", cur_itrs, val_score['Overall Acc'])
                    vis.vis_scalar("[Val] Mean IoU", cur_itrs, val_score['Mean IoU'])
                    vis.vis_table("[Val] Class IoU", val_score['Class IoU'])
                    for k, (img, target, lbl) in enumerate(ret_samples):
                        img = (denorm(img) * 255).astype(np.uint8)
                        target = train_dst.decode_target(target).transpose(2, 0, 1).astype(np.uint8)
                        lbl = train_dst.decode_target(lbl).transpose(2, 0, 1).astype(np.uint8)
                        concat_img = np.concatenate((img, target, lbl), axis=2)  # concat along width
                        vis.vis_image('Sample %d' % k, concat_img)
                model.train()
            scheduler.step()  
            if cur_itrs >=  opts.total_itrs:
                return