# MAE自监督算法在医学图像上的应用
-
MAE自监督算法在医学图像上的应用
本次对上海交通大学的这篇论文进行解析
论文标题: Self Pre-training with Masked Autoencoders for Medical Image Analysis
论文地址: https://arxiv.org/pdf/2203.05573.pdf\摘要
- 将MAE应用在医学图像中
- 由于医学图像没有足够大的数据集,所以作者在预训练和下游任务中使用相同的数据集进行处理
- 为了验证预训练的有效性,在chest X-ray disease 分类,CT abdomen 多器官分割以及MRI brain tumor分割
任务中进行了实验验证,结果突出!
算法介绍
整体思路与MAE基本一致, 创新点有限!
Self Pre-training with Masked Autoencoders
分别介绍编码器、解码器和loss函数
编码器
输入医学图像/体素首先分成不重叠的patches,这些patches然后随机分成visible和masked两组
编码器(例如Vit)仅仅对visible这一组patches进行编码操作
论文并提出MAE的这种方法与MIM方法不一样,能够节省计算量并加速预训练
解码器
MAE的解码器输入来自两个部分,一部分是来自编码器的输出tokens(每个patch的特征表示),第二部分是可学习的masked tokens
通过给所有的tokens添加位置编码,这使得解码器能够识别出masked tokens
为了减小预训练预算,解码器被设计得很轻量化,这里需要注意,解码器只在预训练过程中被使用,不用在下游任务中loss函数
MAE是通过重建函数进行训练的,即mean squares error,是通过预测masked patches中每个像素点/体素的数值计算的
实际操作中,对每个patches进行normalized更有助于重建
Architectures for Downstream Tasks
移除预训练过程中的解码器,使用新的task-specific heads进行下游训练,编码器保持不变
使用线性层进行分类任务
对从Vit输出的class token向量使用线性分类器进行类别预测。
chest X-ray是个多标签任务,所以使用binary cross entropy loss
使用UNETR进行分割任务
因为UNETR也是使用Vit作为编码器的,所以很适合作为下游任务
UNETR的编码器使用预训练的参数权重,解码器随机初始化
实验
实验同样也分为两个部分:预训练和下游任务
编码器首先参与预训练过程以获得良好的训练权重
并直接作为下游任务的初始权重进行监督训练微调
实验细节
backbone:ViT-B/16
优化器:AdamW
patch size:16×16 for 2D图像,16×16×16 for 3D图像
数据预处理增强操作
for Chest X-ray14:
- histogram equalization
- flip and crop \
for BTCV:
- 将数值裁剪在(-175, 250)之间并进行归一化
- random flip and crop \
for BRATS:
- instance-wise normalization
- random flip and crop \
预训练
学习率: 1.5e-4
weight decay: 0.05
下游任务微调
a strategy of layer-wise learning rate decay
random DropPath with a 10% probility
实验结果
对比实验
预训练过程的重建结果
不同数据集实验结果
消融实验
比较不同mask ratio的实验结果