图像分割实战-系列教程4：unet医学细胞分割实战2（医学数据集、图像分割、语义分割、unet网络、代码逐行解读）

杨卓越 Alex Rebel. Yang

2194人浏览 · 2023-12-31 02:33:25

杨卓越 Alex Rebel. Yang · 2023-12-31 02:33:25 发布

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

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

unet医学细胞分割实战1
unet医学细胞分割实战2
unet医学细胞分割实战3
unet医学细胞分割实战4
unet医学细胞分割实战5
unet医学细胞分割实战6

4、train.py主函数解析

4.1 读取配置文件

def main():
    config = vars(parse_args())
    if config['name'] is None:
        if config['deep_supervision']:
            config['name'] = '%s_%s_wDS' % (config['dataset'], config['arch'])
        else:
            config['name'] = '%s_%s_woDS' % (config['dataset'], config['arch'])
    os.makedirs('models/%s' % config['name'], exist_ok=True)
    print('-' * 20)
    for key in config:
        print('%s: %s' % (key, config[key]))
    print('-' * 20)
    with open('models/%s/config.yml' % config['name'], 'w') as f:
        yaml.dump(config, f)

main函数
解析命令行参数为字典
检查 config[‘name’] 是否为 None，如果是
它根据 config[‘deep_supervision’] 的布尔值来设置 config[‘name’]，如果config[‘deep_supervision’] 的值为True
config[‘dataset’] 和 config[‘arch’] 的值，并在末尾添加 ‘_wDS’（表示“with Deep Supervision”）
如果为False，末尾则添加 ‘_woDS’（表示“without Deep Supervision”）
使用 config[‘name’] 来创建一个新目录。这个目录位于 ‘models/’ 目录下，目录名是 config[‘name’] 的值，exist_ok=True 参数的意思是如果目录已经存在，则不会抛出错误
打印符号
打印所有配置参数的名字和默认值
打印符号
根据模型名称创建一个.yaml文件
把所有配置信息全部写入文件中

4.2 定义模型参数

if config['loss'] == 'BCEWithLogitsLoss':
    criterion = nn.BCEWithLogitsLoss().cuda()#WithLogits 就是先将输出结果经过sigmoid再交叉熵
else:
    criterion = losses.__dict__[config['loss']]().cuda()
cudnn.benchmark = True
print("=> creating model %s" % config['arch'])
model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision'])
model = model.cuda()
params = filter(lambda p: p.requires_grad, model.parameters())

定义损失函数，如果损失函数的配置的默认字符参数为BCEWithLogitsLoss
那么使用 PyTorch 中的 nn.BCEWithLogitsLoss 作为损失函数，并且将损失函数的计算移入到GPU中计算，加快速度
如果不是
则从 losses.__dict__ 中查找对应的损失函数，同样使用 .cuda() 方法将损失函数移动到 GPU。（losses.__dict__ 应该是一个包含了多种损失函数的字典，其中键是损失函数的名称，值是相应的损失函数类，这个类是我们自己写的，在后面会解析）
启用 CUDA 深度神经网络（cuDNN）的自动调优器，当设置为 True 时，cuDNN 会自动寻找最适合当前配置的算法来优化运行效率，这在使用固定尺寸的输入数据时往往可以加快训练速度
打印当前创建的模型的名字
动态实例化一个模型， archs 是一个包含多个网络架构的模块， archs.__dict__[config['arch']] 这部分代码通过查找 archs 对象的 __dict__ 属性来动态地选择一个网络架构， __dict__ 是一个包含对象所有属性的字典。在这里，它被用来获取名为 config['arch'] 的网络架构类，config['arch'] 是一个字符串，表示所选用的架构名称
模型放入GPU中

4.3 定义优化器、调度器等参数

if config['optimizer'] == 'Adam':
        optimizer = optim.Adam(
            params, lr=config['lr'], weight_decay=config['weight_decay'])
    elif config['optimizer'] == 'SGD':
        optimizer = optim.SGD(params, lr=config['lr'], momentum=config['momentum'],
                              nesterov=config['nesterov'], weight_decay=config['weight_decay'])
    else:
        raise NotImplementedError

    if config['scheduler'] == 'CosineAnnealingLR':
        scheduler = lr_scheduler.CosineAnnealingLR( optimizer, T_max=config['epochs'], eta_min=config['min_lr'])
    elif config['scheduler'] == 'ReduceLROnPlateau':
        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, factor=config['factor'], patience=config['patience'],
                                                   verbose=1, min_lr=config['min_lr'])
    elif config['scheduler'] == 'MultiStepLR':
        scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[int(e) for e in config['milestones'].split(',')], gamma=config['gamma'])
    elif config['scheduler'] == 'ConstantLR':
        scheduler = None
    else:
        raise NotImplementedError

创建一个过滤器，它筛选出神经网络模型中所有需要梯度（即可训练的）参数， model.parameters()，返回模型的权重和偏置，lambda p: p.requires_grad: 这是一个匿名函数（lambda 函数），用于检查每个参数 p 是否需要梯度
如果优化器是 Adam
则指定参数、学习率、学习率衰减的参数给Adam
如果是SGD
则指定参数、学习率、学习率衰减的参数给SGD，此外还有momentum动量加速，此外还使用了一个自定义的类型转换函数 str2bool 来处理输入值
如果两者都不是
返回错误
如果学习率调度器为CosineAnnealingLR
给该调度器，指定优化器、epochs、最小学习率
如果是ReduceLROnPlateau
给该调度器，指定优化器、指定调整学习率时的乘法因子、指定在性能不再提升时减少学习率要等待多少周期、verbose=1: 这个设置意味着调度器会在每次更新学习率时打印一条信息、最小学习率
如果是MultiStepLR
给该调度器，指定优化器、何时降低学习率的周期数、gamma值
如果是ConstantLR
调度器为None
如果都不是
返回错误

4.4 数据增强

    img_ids = glob(os.path.join('inputs', config['dataset'], 'images', '*' + config['img_ext']))
    img_ids = [os.path.splitext(os.path.basename(p))[0] for p in img_ids]
    train_img_ids, val_img_ids = train_test_split(img_ids, test_size=0.2, random_state=41)
    train_transform = Compose([
        transforms.RandomRotate90(),
        transforms.Flip(),
        OneOf([ transforms.HueSaturationValue(), transforms.RandomBrightness(), transforms.RandomContrast(), ], p=1),
        transforms.Resize(config['input_h'], config['input_w']),
        transforms.Normalize(),
    ])
    val_transform = Compose([
        transforms.Resize(config['input_h'], config['input_w']),
        transforms.Normalize(),
    ])

从本地文件夹inputs，根据config[‘dataset’]的值选择一个数据集，然后images文件，*代表后面所有的文件名称，加上config[‘img_ext’]对应的后缀，返回一个列表，列表的每个元素都是每条数据的路径加文件名和后缀名组成的字符串，类似这种形式：[‘inputs/dataset_name/images/image1.png’, ‘inputs/dataset_name/images/image2.png’, ‘inputs/dataset_name/images/image3.png’]
for p in img_ids按照每个字符串包含的信息，进行遍历，os.path.basename(p)从每个路径 p 中提取文件名，os.path.splitext(...)[0] 则从文件名中去除扩展名，留下文件的基本名称（即 ID）,最后是一个只包含文件名的list，即：[‘image1’, ‘image2’, ‘image3’]
使用sklearn包的train_test_split函数，按照80%和20%的比例分为训练集和验证集，并且打乱数据集，41是随机种子
训练集数据增强
随机以 90 度的倍数旋转图像进行数据增强
水平或垂直翻转图像进行数据增强
从调整色调和饱和度和值（HSV）、随机调整图像的亮度、随机调整图像的对比度这个方式中随机选择一个进行数据增强
将图像调整到指定的高度和宽度
对图像进行标准化（比如减去均值，除以标准差）
验证集同样进行调整，是为了和训练集的尺寸、标准化等保存一致
调整和训练集一样的长宽
调整和训练一样的标准化处理

4.5 数据集制作

 train_dataset = Dataset(
    img_ids=train_img_ids,
    img_dir=os.path.join('inputs', config['dataset'], 'images'),
    mask_dir=os.path.join('inputs', config['dataset'], 'masks'),
    img_ext=config['img_ext'],
    mask_ext=config['mask_ext'],
    num_classes=config['num_classes'],
    transform=train_transform)
val_dataset = Dataset(
    img_ids=val_img_ids,
    img_dir=os.path.join('inputs', config['dataset'], 'images'),
    mask_dir=os.path.join('inputs', config['dataset'], 'masks'),
    img_ext=config['img_ext'],
    mask_ext=config['mask_ext'],
    num_classes=config['num_classes'],
    transform=val_transform)

使用自己写的数据集类制作训练数据集
返回图像数据id
返回图像数据路径
返回掩码数据路径，经过数据增强的时候，albumentations工具包能够自动的把数据做变换的同时标签也相应的变换
返回后缀
返回掩码后缀
分类的种类
数据增强（这里制定为None）
同样的给验证集也来一遍

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=config['batch_size'],
    shuffle=True,
    num_workers=config['num_workers'],
    drop_last=True)
val_loader = torch.utils.data.DataLoader(
    val_dataset,
    batch_size=config['batch_size'],
    shuffle=False,
    num_workers=config['num_workers'],
    drop_last=False)
    log = OrderedDict([ ('epoch', []), ('lr', []), ('loss', []), ('iou', []), ('val_loss', []), ('val_iou', []), ])

制作训练集Dataloader
指定训练数据集
batch_size
洗牌操作
进程数
不能整除的batch是否就不要了
同样的给验证集也来一遍
最后一行日志记录：创建OrderedDict 对象 log，将’epoch’、‘lr’、‘loss’、‘iou’、‘val_loss’、'val_iou’按照类似字典的形式进行存储（与字典不同的是它会记住插入元素的顺序）

4.6 迭代训练

    best_iou = 0
    trigger = 0
    for epoch in range(config['epochs']):
        print('Epoch [%d/%d]' % (epoch, config['epochs']))
        train_log = train(config, train_loader, model, criterion, optimizer)
        val_log = validate(config, val_loader, model, criterion)
        if config['scheduler'] == 'CosineAnnealingLR':
            scheduler.step()
        elif config['scheduler'] == 'ReduceLROnPlateau':
            scheduler.step(val_log['loss'])
        print('loss %.4f - iou %.4f - val_loss %.4f - val_iou %.4f'
              % (train_log['loss'], train_log['iou'], val_log['loss'], val_log['iou']))
        log['epoch'].append(epoch)
        log['lr'].append(config['lr'])
        log['loss'].append(train_log['loss'])
        log['iou'].append(train_log['iou'])
        log['val_loss'].append(val_log['loss'])
        log['val_iou'].append(val_log['iou'])
        pd.DataFrame(log).to_csv('models/%s/log.csv' % config['name'], index=False)
        trigger += 1
        if val_log['iou'] > best_iou:
            torch.save(model.state_dict(), 'models/%s/model.pth' % config['name'])
            best_iou = val_log['iou']
            print("=> saved best model")
            trigger = 0
        if config['early_stopping'] >= 0 and trigger >= config['early_stopping']:
            print("=> early stopping")
            break
        torch.cuda.empty_cache()

记录最好的IOU的值
trigger 是一个计数器，用于追踪自从模型上次改进（即达到更好的验证 IoU）以来经过了多少个训练周期（epochs），这种技术通常用于实现早停（early stopping）机制，以避免过度拟合
按照epochs进行迭代训练
打印当前epochs数，即训练进度
使用训练函数进行单个epoch的训练
使用验证函数进行单个epoch的验证
当使用 CosineAnnealingLR 调度器时
scheduler.step() 被直接调用，无需任何参数
当使用 ReduceLROnPlateau 调度器时
scheduler.step(val_log['loss']) 调用时传入了验证集的损失 val_log[‘loss’] 作为参数
打印当前epoch训练损失、训练iou
打印当前epoch验证损失、验证iou
当前epoch索引加入日志字典中
当前学习率值加入日志字典中
当前训练损失加入日志字典中
当前训练iou加入日志字典中
当前验证损失加入日志字典中
当前验证iou加入日志字典中
当前日志信息保存为csv文件
trigger +1
如果当前验证iou的值比当前记录最佳iou的值要好
保存当前模型文件
更新最佳iou的值
打印保存了当前的最好模型
把trigger 置0
如果当前记录的trigger的值大于提前设置的trigger阈值
打印提前停止
停止训练
清除GPU缓存