kaggle《SIIM-ACR Pneumothorax Segmentation》第一名方案详解

当初刷微博无意中看到这个比赛，内容和自己所在团队的工作内容很契合，就花了点时间学习了下第一名的方案。在这里把自己学习的成果记录下来，也希望和大家一起讨论学习。

文章同步发在我的个人博客，欢迎大佬们指教。kaggle《SIIM-ACR Pneumothorax Segmentation》第一名方案详解

1 Introduction

引用一段比赛页面对这个比赛的介绍：

“Pneumothorax can be caused by a blunt chest injury, damage from underlying lung disease, or most horrifying—it may occur for no obvious reason at all. On some occasions, a collapsed lung can be a life-threatening event.
Pneumothorax is usually diagnosed by a radiologist on a chest x-ray, and can sometimes be very difficult to confirm. An accurate AI algorithm to detect pneumothorax would be useful in a lot of clinical scenarios. AI could be used to triage chest radiographs for priority interpretation, or to provide a more confident diagnosis for non-radiologists.”

比赛地址：SIIM-ACR Pneumothorax Segmentation

Data:

• 10679份dicom文件（站立位胸片），有气胸：无气胸 = 2379:8300

• 有气胸的胸片皆有mask, run-length-encoded (RLE)格式

• 一个胸片如果有多处气胸，会有多个单独的mask

Evaluation: Dice coefficient:

$$\frac{2\ast |X\bigcap Y|}{|X|+|Y|}$$

2 OverView

用一张图来表示作者的思路：

作者在训练模型时，采用5折交叉验证，并将模型训练分成了4个阶段（下文中分别表示成part0，part1，part2，part3），每个阶段都是上图中这样一个完整的过程，后一阶段直接在前一阶段的最优模型上fine-tuning.每一个阶段有不同的数据构成和模型参数。

3 Input Data

3.1 Data Augmentation

这里模型的输入是1024X1024X3的胸片和1024X1024X1的mask。作者直接使用了图像增强库 albumentations 对数据进行增强，作者使用了一个较为复杂的方案：以固定顺序对图像进行不同的变换，并且给予每种方法一定的概率，使增强方法的运用随机化。具体如下:

1. HorizontalFlip：0.5
2. OneOf：0.3
   1. RandomContrast：0.5
   2. RandomGamma：0.5
   3. RandomBrightness：0.5
3. OneOf：0.3
   1. ElasticTransform：0.5
   2. GridDistortion：0.5
   3. OpticalDistortion：0.5
4. ShiftScaleRotate：0.5

上面小数表示此增强方法运用的概率，OneOf表示在其子方法中选择一个。

需要注意的是，这里有2个OneOf，第一个OneOf下面的增强方法主要对图像的亮度、对比度等进行调整，而第二个OneOf下面的增强方法主要对图像的形状进行调整。可以看几个例子。

因为训练模型分成四个阶段，每个阶段使用的数据是一样的，使用一定概率给图像做增强，实际上保证了每个阶段实际参加训练的数据都不完全一样。

3.2 Sliding Sample

对无气胸的样本随机下采样，使有气胸的样本占0.8（part0），0.6（part1），0.4（part2），0.5（part3）

4 Model

4.1 Model Zoo

• Albunet: UNet (论文链接) with Resnet34(论文链接) encoder proposed by Alexander Buslaev

UNet简直是图像分割的神器，特别是在医疗图像分割的上，目前其各种变体网络仍是各种比赛的主力。Unet主要有2个特点：1.U型结构；2.skip connection.

– U型结构encoder的下采样和decoder的上采样的次数相同，这就保证了模型的输出恢复到原图片的分辨率，实现端到端的预测。

– skip connection的结构使模型结合不同level的feature map上进行学习，相比于FCN分割边缘更清晰。

同时Albunet的网络的encoder使用的是ResNet，其由何凯明大佬于2015年提出（Unet也是这一年提出，回过头看，这一年真的是丰收的一年），同样风靡至今。从理论上讲，越高级的特征，应该有越强的表征能力，而VGG网络证明，网络的深度对特征的表达能力至关重要。

理想情况下，当我们直接对网络进行简单的堆叠到特别长，网络内部的特征在其中某一层已经达到了最佳的情况，这时候剩下层应该不对改特征做任何改变，自动学成恒等映射的形式。也就是说，对一个特别深的深度网络而言，该网络的浅层形式的解空间应该是这个深度网络解空间的子集，但实际上，如果使用简单是网络堆叠，由于网络性能衰减，网络的效果反而越差。为了深度网络后面的层至少实现恒等映射的作用，作者提出了residual模块。

除了Albunet，作者还尝试了如下2种网络进行实验：

• Resnet50: GitHub链接

• SCSEUnet：也就是SENet((论文：squeeze-and-excitation network))

4.2 Loss Function

作者使用BCE，Dice和Focal loss加权的方式作为最终的损失函数：

$$Loss=W_1\ast StableBCELoss+W_2\ast DiceLoss+W_3\ast FocalLoss2d()$$

这里需要注意的是，损失函数StableBCELoss的输入是没有进行sigmoid计算的模型输出，而其他2个损失函数的输入是经过sigmoid计算的模型输出。作者使用这三个损失函数，是想从三个维度驱动网络学习。

– Focal loss主要解决正负样本失衡的问题，在医学图像分割的这种像素级分类的任务中，往往都是正样本较少，负样本较多，气胸的分割同样不例外。其公式如下：

从如上公式上可以看到，某个样本输出的概率越高，其产生的loss越小，实际上达到了促使网络学习困难样本的目的（在分割任务中，目标像素就是困难样本，因为其数量较少）。

– Dice loss在很多关于医学图像分割的竞赛、论文和项目中出现的频率很高，此比赛的评价指标也是它。从前文它的公式可以看到，其表示的是预测的轮廓与真实的mask的相似程度。其公式如下：

$$DiceLoss=1-\frac{2\ast |X\bigcap Y|}{|X|+|Y|}$$

– StableBCELoss, 这里非常奇怪的是，其输入是没有进行sigmoid计算的模型输出，至今没想通（to do：等搞明白了补上）。这里附上pytorch文档中其公式：

最终作者寻优，得到各个损失函数的权重如下：

• (3,1,4) for albunet_valid and seunet;

• (1,1,1) for albunet_public;

• (2,1,2) for resnet50.

4.3 Learn Rate Scheduler

作者在不同的训练阶段，采用了不同的学习率调整策略，具体如下：

Phase	Start Lr	Scheduler
part0	1e-4	ReduceLROnPlateau
part1	1e-5	ReduceLROnPlateau or CosineAnnealingWarmRestarts
part2	1e-5	ReduceLROnPlateau or CosineAnnealingWarmRestarts
part3	1e-6	ReduceLROnPlateauor CosineAnnealingWarmRestarts

4.3.1 torch.optim.lr_scheduler.ReduceLROnPlateau

其基本思想是：当参考的评价指标停止变优时,降低学习率，挺实用的方法。

Code Example:

optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = ReduceLROnPlateau(optimizer, 'min')
for epoch in range(10):
     train(...)
     val_loss = validate(...)
     # Note that step should be called after validate()
     scheduler.step(val_loss)

4.3.2 torch.optim.lr_scheduler.CosineAnnealingLR

可以翻译为余弦退火调整学习率：根据论文SGDR: Stochastic Gradient Descent with Warm Restarts而来，只包含论文里余弦退火部分，并不包含restart部分。

当使用SGD时，模型的Loss应越来约接近全局最小值。当它逐渐接近这个最小值时，学习率应该变得更小来使得模型尽可能接近这一点。

从上图可以看出，随着x的增加，余弦值首先缓慢下降，然后加速下降，再次缓慢下降。这种下降模式能和学习率配合，比较有效的训练模型。

4.3.3 torch.optim.lr_scheduler.CosineAnnealingWarmRestarts

SGDR：这个方法就是论文SGDR: Stochastic Gradient Descent with Warm Restarts的完整实现。

在训练时，梯度下降算法可能陷入局部最小值，而不是全局最小值。做随机梯度下降时可以通过突然提高学习率，来“跳出”局部最小值并找到通向全局最小值的路径。具体可以看下图：

5 Post Proprecess

我们训练的分割模型输出每个像素的0-1概率，然后卡一下阈值，我们可以称这样的mask为basic sigmoid mask, 实际上医学图像中我们的分割目标也许并不存在，所以常用双重阈值（top_score_threshold, min_contour_area）的方法计算出mask并同时判断是否有分割目标（在本次比赛中我们分割目标是气胸），这种方法且称为doublet。其具体逻辑为：当大于概率阈值top_score_threshold的pixel数少于
min_contour_area，就将mask像素值全部置0，也就是认为此胸片没有气胸。简单画个示意图如下：

而作者在此基础上作了改进，使用了三重阈值（top_score_threshold, min_contour_area, bottom_score_threshold）的方法来达到相同的目标，且称改进后的方法为Triplet.其具体逻辑为：当大于概率阈值top_score_threshold的pixel数少于
min_contour_area，就将此mask pixel值全部置0，也就是认为此胸片没有气胸，然后再使用阈值bottom_score_threshold产生真正的mask。简单画个示意图如下：

最终作者通过搜索，分别获得了在validation和在Public Leaderboard上的最优参数：

• Best triplet on validation: (0.75, 2000, 0.3).
• Best triplet on Public Leaderboard: (0.7, 600, 0.3)

最后再附上作者的代码：

classification_mask = predicted > top_score_threshold
mask = predicted.copy()
mask[classification_mask.sum(axis=(1,2,3)) < min_contour_area, :,:,:] = np.zeros_like(predicted[0])
mask = mask > bot_score_threshold
return mask

6 Train and Inference

6.1 Train

前面就已经提到过，作者将训练过程分为了4个阶段，这里在最前面新加一个预训练阶段。梳理如下：

Phase	Sample rate	Start lr	Scheduler	Epochs	Note
part pre					The model be pretrained on our dataset with lower resolution (512x512)
part0	0.8	1e-4	ReduceLROnPlateau	10-12	The goal of this part: quickly get a good enough model with validation score about 0.835
part1	0.6	1e-5	ReduceLROnPlateau or CosineAnnealingWarmRestarts	Repeat until best convergence	uptrain the best model from the previous step
part2	0.4	1e-5	ReduceLROnPlateau or CosineAnnealingWarmRestarts	Repeat until best convergence	uptrain the best model from the previous step
part3	0.5	1e-6	ReduceLROnPlateauor CosineAnnealingWarmRestarts	Repeat until best convergence	uptrain the best model from the previous step

6.2 Inference

作者采用五折交叉验证训练模型，并选择每一个fold的top3个模型的结果求平均输出最终的mask。这里的求平均每个像素点的概率求平均。模型的最终效果如下表：

NOTE：

• albunet_public - best model for Public Leaderboard

• albunet_valid - best resnet34 model on validation

• seunet - best seresnext50 model on validation

• resnet50 - best resnet50 model on validation