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

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

有任何问题欢迎在下面留言
本篇文章的代码运行界面均在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)

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

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())

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

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

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(),])

1. 从本地文件夹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’]

2. for p in img_ids按照每个字符串包含的信息，进行遍历，os.path.basename(p)从每个路径 p 中提取文件名，os.path.splitext(...)[0] 则从文件名中去除扩展名，留下文件的基本名称（即 ID）,最后是一个只包含文件名的list，即：[‘image1’, ‘image2’, ‘image3’]

3. 使用sklearn包的train_test_split函数，按照80%和20%的比例分为训练集和验证集，并且打乱数据集，41是随机种子

4. 训练集数据增强

5. 随机以 90 度的倍数旋转图像进行数据增强

6. 水平或垂直翻转图像进行数据增强

7. 从调整色调和饱和度和值（HSV）、随机调整图像的亮度、随机调整图像的对比度这个方式中随机选择一个进行数据增强

8. 将图像调整到指定的高度和宽度

9. 对图像进行标准化（比如减去均值，除以标准差）

10. 验证集同样进行调整，是为了和训练集的尺寸、标准化等保存一致

11. 调整和训练集一样的长宽

12. 调整和训练一样的 标准化处理

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)

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

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', []), ])

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

4.6 迭代训练

    best_iou = 0trigger = 0for 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 += 1if 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 = 0if config['early_stopping'] >= 0 and trigger >= config['early_stopping']:print("=> early stopping")breaktorch.cuda.empty_cache()

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

自此，train.py的main函数部分全部解读完毕，其中有多个子函数或者类，在下一篇文章中继续解读

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