[论文解读] U-Net-based Models for Skin Lesion Segmentation: More Attention and Augmentation

Pooya Mohammadi Kazaj, MohammadHossein Koosheshi+2 more|arXiv (Cornell University)|Oct 28, 2022

Cutaneous Melanoma Detection and Management被引用 6

一句话总结

本研究提出改进的U-Net架构，集成CBAM与金字塔注意力模块，并结合先进的数据增强技术（包括自定义头发增强及CutMix/Mosaic），以提升ISIC 2016数据集上的皮肤病变分割性能。在特定数据增强设置下，U-Net-ResNet50与R2U-Net模型达到最高的Dice与IoU分数，所提出的U-Net Pyramid CBAM Gate在处理类内与类间不一致性方面显著优于其他模型。

ABSTRACT

According to WHO[1], since the 1970s, diagnosis of melanoma skin cancer has been more frequent. However, if detected early, the 5-year survival rate for melanoma can increase to 99 percent. In this regard, skin lesion segmentation can be pivotal in monitoring and treatment planning. In this work, ten models and four augmentation configurations are trained on the ISIC 2016 dataset. The performance and overfitting are compared utilizing five metrics. Our results show that the U-Net-Resnet50 and the R2U-Net have the highest metrics value, along with two data augmentation scenarios. We also investigate CBAM and AG blocks in the U-Net architecture, which enhances segmentation performance at a meager computational cost. In addition, we propose using pyramid, AG, and CBAM blocks in a sequence, which significantly surpasses the results of using the two individually. Finally, our experiments show that models that have exploited attention modules successfully overcome common skin lesion segmentation problems. Lastly, in the spirit of reproducible research, we implement models and codes publicly available.

研究动机与目标

• 通过注意力机制与数据增强技术提升深度学习模型在皮肤病变分割任务中的性能。
• 解决常见分割挑战，如类内不一致性（例如病变区域的颜色差异）与类间模糊性（例如正常皮肤被误判为病变）。
• 通过新提出的指标ΔM（训练与验证指标之差）评估模型泛化能力与过拟合程度。
• 开发并验证一种新型注意力模块——U-Net Pyramid CBAM Gate，其在复杂场景下优于单一注意力模块。
• 通过公开发布代码与训练好的模型，确保研究的可复现性。

提出的方法

• 在跳跃连接中集成CBAM与AG注意力模块的新型U-Net-ResNet50与R2U-Net模型，以增强特征注意力。
• 设计一种新型序列注意力模块，结合金字塔、AG与CBAM模块，以提升特征表示能力与分割鲁棒性。
• 应用四种数据增强配置，包括基于Albumentation的增强、Mosaic、CutMix，以及两种自定义头发增强技术，以减少模型因头发伪影产生的混淆。
• 在ISIC 2016数据集上使用五折交叉验证训练十个基于U-Net的模型，并采用Dice、IoU与Focal Tversky指标评估性能。
• 引入ΔM作为新指标，通过测量各训练周期中训练与验证分数的差距，量化过拟合程度。
• 采用基于每轮训练平均训练周期数的新型训练速度评估方法，以比较模型效率。

实验结果

研究问题

• RQ1 CBAM与AG等注意力模块如何提升基于U-Net的模型在皮肤病变分割中的性能？
• RQ2 将多种注意力模块（金字塔、AG、CBAM）组合使用，相较于单独使用，其影响如何？
• RQ3 哪些数据增强策略——尤其是自定义头发增强与CutMix/Mosaic——最有效地减少过拟合并提升泛化能力？
• RQ4 在分割精度、训练速度及对类内与类间不一致性鲁棒性方面，U-Net-ResNet50与R2U-Net模型表现如何比较？
• RQ5 注意力模块在多大程度上缓解了常见分割问题，如边界检测不佳与因颜色相似导致的误分类？

主要发现

• U-Net-ResNet50在使用Aug-1与Aug-2数据增强配置下，达到最高Dice分数（0.91±0.0）与IoU分数（0.83±0.0），表明其分割性能最优。
• R2U-Net展现出最快的训练速度（平均7.8个训练周期），适用于计算资源有限的环境。
• U-Net Pyramid CBAM Gate模型在涉及类间模糊性与类内不一致性的场景中，显著优于U-Net-ResNet50、Double U-Net与R2U-Net，尤其在颜色变化区域表现更优。
• CBAM模块显著提升了R2U-Net在边界检测与浅色区域的鲁棒性，有效降低了类内不一致性。
• R2U-Net与U-Net-ResNet50模型的过拟合程度最低（Dice指标的ΔM分别为0.009与0.015），表明其泛化能力出色。
• 自定义头发增强与CutMix/Mosaic显著减少了模型对头发伪影的混淆，并提升了对形态变化的鲁棒性。