冠军方案解读 | nnUNet改进提升笔记

点击上方“计算机视觉工坊”，选择“星标”

干货第一时间送达

作者丨ChaucerG

来源丨集智书童

本文描述了基于nn-UNet试验了几种改进，包括使用更大的网络、用GN替换BN以及在解码器中使用Axial Attention。与Baseline相比，量化指标略有改进。在unseen test data的最终排名中，nn-UNet赢得了第一名的好成绩。

1改进策略

1.1 数据方面

BraTS2021包括了来自2000名患者的多参数MRI扫描结果，其中1251人的图像提供了分割标签给参与者来开发算法，其中219人在验证阶段被用于公共排行榜，其余530个案例用于私人排行榜和参与者的最终排名。

MRI扫描有4种对比:

• 原生T1加权图像

• 对比后T1加权(T1GD)

• T2加权

• T2流体衰减反转恢复(T2-Flair)

注释由1-4个评分员手工完成，最终得到有经验的神经放射学家的批准。这些标签包括gd增强肿瘤(ET)、瘤周水肿/侵袭组织(ED)和坏死肿瘤核心(NCR)区域。所有MRI扫描均通过同解剖模板配准、各向同性1mm3分辨率插值和颅骨剥离进行预处理。所有MRI扫描及相关标记的图像大小为240×240×155。

图1

图1显示了这4种对比与分割的代表性切片。在输入网络之前，对提供的数据进行了进一步的处理。为了减少计算量，将volumes裁剪为non-zero voxels。由于MR图像的强度是定性的，因此根据其均值和标准差对voxels进行归一化。

1.2 模型方面

1、Baseline nnU-Net

nnU-Net核心是一个在128×128×128大小的Patch上运行的3D U-Net。该网络具有编码-解码器结构，并带有Skip Connection，将两条路径连接起来。

该编码器由5个相同分辨率的卷积层组成，具有卷积下采样功能。该解码器遵循相同的结构，使用转置卷积上采样和卷积操作在同一级别上的编码器分支的串联Skip特征。每次卷积操作后，采用斜率为0.01的Leaky ReLU(lReLU)和批归一化处理。mpMRI volumes被连接并作为4通道输入。

nnU-Net应用Region-Based训练，而不是预测3个相互排斥的肿瘤子区域，而不是预测3个互斥肿瘤分区,与提供的分割标签一样，该网络预测的是增强肿瘤的3个重叠区域如加强肿瘤(ET,original region),肿瘤核心或TC(ET+necrotic tumor),和整个肿瘤或WT(ET+NT+ED)。

网络的最后一层的softmax被sigmoid所取代，将每个voxels作为一个多类分类问题。

由于公共和私人排行榜的计算指标是基于这些区域的，这种基于区域的训练可以提高表现。额外的sigmoid输出添加到每个分辨率除了2个最低的水平，应用深度监督和改善梯度传播到早期层。卷积滤波器的数量被初始化为32个，并且分辨率每降低一倍，最大可达320个。

2、更大的网络和GN

第一个修改是，通过将编码器中的kernel数量加倍，同时在解码器中保持相同的kernel，非对称地增加了网络的大小。由于训练数据的数量是前一年的4倍，增加网络的容量将有助于它能够建模更大的数据种类。kernel的最大数量也增加到512个，改进后的网络结构如图2所示：

图2

第二个修改是用GN代替所有BN。即使使用混合精度训练，3D卷积网络也需要大量的GPU内存，这限制了在训练中可以使用的Batch-Size。

3、Axial attention解码器

最后添加的是在解码器中使用Axial attention。Self-Attention或Transformer是一个突破性的想法，允许学习一个输入序列的自适应注意力仅仅基于它自己。Self-Attention最初是在自然语言处理中，现在已经慢慢被计算机视觉研究所采用。当试图将Self-Attention应用于视觉问题时，主要的障碍之一是注意力机制的计算复杂度与输入的大小成二次方，这使得它不可能适合或训练网络在一个标准的工作站设置。当处理带有额外维度的3D数据时，这是一个更大的问题。

Axial attention最近被提出作为将注意力应用于多维数据时的一种有效解决方案。通过将Self-Attention独立地应用于输入的每一个轴上，计算只与图像大小成线性比例，使注意力机制即使与3D数据整合成为可能。

本文将Axial attention应用到网络的解码器上，将其运行在转置卷积上采样的输出上，然后将它们相加。

图3

显示了Axial attention解码器块的示意图。即使有了更有效的注意力，作者发现这种方法也不可能应用于最高分辨率的特征(128×128×128)，因此只选择了4个较低分辨率的特征。注意力头的数量和每个头的尺寸随着分辨率的降低而增加一倍，分别从4和16(64×64×64分辨率)开始。

1.3 训练策略

这里遵循nnU-Net训练方法。每个网络都接受了5倍交叉验证的训练。在训练过程中，动态地应用数据增强来提高泛化能力。数据增强包括随机旋转和缩放、弹性变形、附加亮度增强和伽玛缩放。

优化的目标是binary entropy loss和Dice loss的总和，计算在最终的全分辨率输出以及在低分辨率的辅助输出。使用Batch Dice loss代替sample Dice loss，将整个批次作为一个样本来计算损失，而不是平均每个样本在小批中的Dice。批量Dice帮助稳定训练通过减少来自样本的错误，少量的注释样本。网络采用Nesterov动量为0.99的随机梯度下降法进行优化。初始学习率为0.01，并按照多项式计划衰减：

每次训练运行持续1000个epoch，每个epoch包含250个小批量。当前折叠验证集上的Dice Score用于监视训练进度。所有实验都是在24GB VRAM的NVIDIA RTX 3090 GPU上使用Pytorch 1.9进行的。开发了以下模型:

• BL：Baseline nnUNet

• BL+L：Baseline with Large nnUNet

• BL+GN：Baseline with Group Normalization

• BL+AA：Baseline with axial attention, batch normalization

• BL+L+GN：nnUNet with larger Unet, group normalization

2比赛结果

3参考阅读

[1].Extending nn-UNet for brain tumor segmentation

本文仅做学术分享，如有侵权，请联系删文。

重磅！计算机视觉工坊-学习交流群已成立

扫码添加小助手微信，可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群，旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。

同时也可申请加入我们的细分方向交流群，目前主要有ORB-SLAM系列源码学习、3D视觉、CV&深度学习、SLAM、三维重建、点云后处理、自动驾驶、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、深度估计、学术交流、求职交流等微信群，请扫描下面微信号加群，备注：”研究方向+学校/公司+昵称“，例如：”3D视觉 + 上海交大 + 静静“。请按照格式备注，否则不予通过。添加成功后会根据研究方向邀请进去相关微信群。原创投稿也请联系。

▲长按加微信群或投稿

▲长按关注公众号

3D视觉从入门到精通知识星球：针对3D视觉领域的视频课程（三维重建系列、三维点云系列、结构光系列、手眼标定、相机标定、激光/视觉SLAM、自动驾驶等）、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答五个方面进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业为一体的铁杆粉丝聚集区，近4000星球成员为创造更好的AI世界共同进步，知识星球入口：

学习3D视觉核心技术，扫描查看介绍，3天内无条件退款

 圈里有高质量教程资料、可答疑解惑、助你高效解决问题

觉得有用，麻烦给个赞和在看~