Tensorflow入门教程（五十一）——H2NF-Net

今天将分享Unet的改进模型H2NF-Net，改进模型来自2020年的论文《H2NF-Net for Brain Tumor Segmentation using Multimodal MR Imaging: 2nd Place Solution to BraTS Challenge 2020 Segmentation Task》，简单明了给大家分析理解该模型思想。

一、H2NF-Net网络优点

论文提出了一种高分辨率的非局部特征网络（HNF-Net）来对多模式MR图像进行脑肿瘤分割。HNF-Net主要是基于并行多尺度融合（PMF）模块来构建，该模块具有强大的高分辨率特征表示，并能聚集多尺度上下文信息。期望最大化注意（EMA）模块也被引入到该模型中，以可接受的计算复杂性为代价来增强远距离上空间上下文信息。

论文还进一步提出了一种混合高分辨率和非局部特征网络（H2NF-Net）来解决这一具有脑肿瘤分割挑战性的任务。与原始的HNF-Net相比，拟议的H2NF-Net添加了两级级联的HNF-Net，并使用单个模型和级联模型对不同的脑肿瘤子区域进行了精细分割。

二、H2NF-Net网络

在介绍的H2NF-Net结构前，先对HNF-Net及其两个关键模块进行简要介绍。还提供了两级级联HNF-Net的详细信息。

2.1、单HNF-Net网络

单个HNF-Net具有编码器-解码器结构。首先将四个MR序列连接起来，形成一个四通道输入，然后进行五个尺度分析处理，即r，1 /2r，... 1 / 16r，以绿色，黄色，蓝色，粉红色高亮显示，在原始比例r处，有四个卷积块，两个用于编码，另外两个用于解码。从编码器到解码器的连接会跳过其他尺度的处理，在通过长范围残差融合以保持高分辨率和高空间信息。在其他四个尺度上，四个PMF模块一起用作高分辨率和多尺度聚合特征提取器。在最后一个PMF模块的末尾，先将四个尺度的输出特征图恢复到1 / 2r比例，然后将其拼接为混合特征图。接下来，EMA模块用于有效地捕获长范围相关的上下文信息，并减少所获得的冗余混合特征。最后，通过1x 1x 1卷积和上采样将EMA模块的输出恢复到原始尺度r和32通道，然后将其添加到由编码器生成的全分辨率特征图上，以密集预测体素标签。所有的下采样操作均采用步长为2的卷积操作完成的，所有上采样操作均通过联合1x 1x 1卷积和三线性插值实现。

PMF模块。PMF模块被设计为两部分，即并行多尺度卷积块和全连接融合块。前者具有类似于组卷积的一组并行分支，并且每个分支都以特定比例使用重复的残差卷积块构建。后者以并行但全连接的方式融合了并行多尺度卷积块的所有输出特征，其中每个分支是所有分辨率分支的输出特征的总和。因此，在每个PMF模块中，并行多尺度卷积块可以充分利用多分辨率特征，但保留高分辨率特征表示，并且全连接的融合块可以聚合丰富的多尺度上下文信息。此外，级联多个PMF模块，其中分支的数量随深度而逐渐增加，如图2（a）所示。从最高分辨率阶段的角度来看，它的高分辨率特征表示通过多尺度低分辨率表示的重复融合得到了增强。同时，级联PMF模块可以看作是多个深度和宽度不同的U形子网的集合，这往往会进一步缩小不同深度的特征的语义鸿沟。

EMA模块。事实证明，非局部的自我注意机制可以帮助聚集来自所有空间位置的上下文信息并捕获长距离的依赖关系。但是，潜在的高计算复杂性使其很难应用于3D医学图像分割任务。因此，将EMA模块引入我们的肿瘤分割模型，旨在将轻量级的非局部注意机制整合到的模型中。EMA模块的主要概念是在一组特征重建基础上进行非局部注意，而不是直接在高分辨率特征图上实现。由于重建基础的元素比原始特征图少得多，因此可以显着减少非局部注意力的计算成本。如图2（b）的下部所示，输入特征图X的形状为CxDxHxW，其中C为通道编号。虚线框表示EM算法的操作， 其中，非局部空间注意图A和基础u分别通过E步和M步来交替生成。收敛后，我们可以使用获得的A和u来生成重建的特征图X~。在EM算法操作的开始和结束时，均采用1x1x1卷积来更改通道数。另外，为了避免过度拟合，进一步将X~与X求和进行残差融合。

4.2、级联HNF-Net

进一步构建了两级层叠的HNF-Net，其结构如图3所示。与单个HNF-Net相比，级联的第一级网络HNF-Net较窄，没有EMA模块。第二个网络具有与单个HNF-Net相同的结构，但是接收原始图像块的串联和第一个网络的预测作为输入。两阶段级联的HNF-Net以端到端的方式进行训练，并且还为第一阶段网络的输出添加了深度监控，以进行稳定的训练。

3、实验设置与结果

3.1、数据

BraTS 2020挑战数据集包含369个训练，125个验证和166个测试多模式脑MR研究。每个研究都有四个MR序列，包括T1加权（T1），对比度T1加权（T1ce），T2加权（T2）和uid衰减反转恢复（Flair）序列，如图1所示。所有 MR图像的大小相同，为240 240 155，体素间距为1 x1x 1mm3。对于每项研究，增强肿瘤（ET），肿瘤周围水肿 （ED），坏死和非增强型肿瘤核心（NCR / NET）由专家逐个像素进行注释。 训练注释数据是公开可用的，而验证和测试数据的注释则分别保留用于在线评估和最终分类比赛。

3.2、实现细节

预处理。对每个脑部MR序列独立进行了一套预处理，包括脑部剥离，以[0.5％-99.5％]的窗口剪切所有脑部体素强度，并将其标准化为零均值 和单位方差。

训练。在训练阶段，将输入图像随机裁剪为固定的128x128x128大小，并沿着通道维度将四个MR序列连接起来作为模型的输入。训练迭代设置为450个epoch。使用Adam优化器对模型进行了训练，批次大小为4，beta值为（0.9，0.999）。初始学习率设置为0.0085,然后乘以（1-current epoch/max epoch）的0.9次方。还以1e-5的l2权重衰减使训练正规化。为了减少潜在的过拟合现象，进一步采用了几种在线数据增强方法，包括随机倾斜（在所有三个平面上均独立），随机旋转（在所有三个平面上均正负10度），每通道随机强度[0.1] ，强度缩放比例为[0.9-1.1]。在所有实验中，根据经验设置基数K = 256。采用广义骰子损失和二进制交叉熵损失的组合作为损失函数。所有实验均基于PyTorch 1.2.0进行。由于级联HNF-Net的训练需要超过11 Gb的GPU内存，因此使用了4个NVIDIA Tesla P40 GPU和4个NVIDIA Geforce GTX 2080Ti GPU来训练级联 HNF-Net和单个HNF-Net。

推理。在推论阶段，首先裁剪原始图像，尺寸为224x160x155，这是基于对整个数据集的统计分析确定的，以覆盖整个大脑区域，但是具有最小的冗余背景体素。然后，用滑动块分割裁剪的图像，而不是立即预测整个图像，输入块的大小和滑动步幅分别设置为128x128x128和32x32x27。对于每个推理patch，采用了测试时间增强（TTA）来进一步提高分割性能，其中包括7中不同翻转（（x），（y），（z），（x，y），（x，z），（y，z），（x，y，z），其中x，y，z分别表示三个轴）。然后，对增强和部分重叠patch块的预测取平均，以生成整个图像分割结果。最后通过替换增强型肿瘤进行了后处理，当预测的增强肿瘤的体积小于阈值时，使用NCR / NET。根据验证集上的结果，对于单个模型和级联模型，阈值分别设置为300和500。

3.3、结果

H2NF-Net对测试集进行了分割，获得了平均Dice分数0.78751、0.91290和0.85461以及Hausdorff距离（95 %ET，WT和TC分别为26.57525、4.18422和4.97162），在BraTS 2020挑战赛分割任务的近80名参与者中赢得了第二名。图4可视化了测试集的某些扫描结果的分割结果。