文章概要
文章先是在开头花了一半的篇幅介绍了2D丢失了层间信息、3D训练受制于设备和LSTM的困难。然后介绍他们的AH-Net可以从二维图像中学习到的卷积特征转化为三维各向异性体,并在两个数据及上表现出了state-of-the-art的成绩。
Introduction
直接利用3D信息存在着一些困难,主要体现在如下几个方面:
- 较小的3×3×3核可能很难从三维信息中学习有用的特征,因为每个维度的信息密度不同
- 理论上,与二维网络相比,三维网络需要更多的特征。三维网络的性能受到GPU内存的限制,限制了网络的宽度和深度
- 不像现在可以利用数百万张2D图像预训练,3D任务大多需要从头开始训练,此缺乏大型的三维数据集。
AH-Net
AH-Net使用预先训练的权重初始化,并使用多个2D图像切片作为输入,类似于伪RGB图片,组成多通道。该二维网络的特征编码器通过对二维核进行一维扩展,将其转化为三维网络,然后加入一个特征解码器子网络来提取三维上下文,卷积核有3x3x1和1x1x3两种。在输出之前,经过一个金字塔池化模块(详见SPP-Net)来提取多尺度特征。
由于是2D的模型,所以AH-Net也可以利用ImageNet的庞大数据集进行预训练,这样既易于训练,又易于推广。
AH-Net由一个特征编码器和一个特征解码器组成.从二维网络转换而来的编码器,是为了从高分辨率的二维切片中提取深度表示而设计的。由各向异性卷积的密集连接块构成的译码器负责利用3D上下文并保持两层间的一致性。
网络训练分两个阶段进行:学习编码器;然后添加3D解码器,并在编码器参数锁定的情况下进行精细调整。这里需要说明的是文中LOSS函数用的是Focal Loss,其之前一直被用在目标检测中。
Learning a multi-channel 2D feature encoder
作者训练了一个类似于 Large Kernel
Matters – Improve Semantic Segmentation by Global Convolutional Network 中提出的结构的2D多通道全局卷积网络(mc-gcn)来提取不同分辨率的二维片内特征。主要框架是ResNet50,采用了ImageNet数据集作为预训练,然后用从3d卷中提取的2D图像切片对网络进行精细调整。这个网络的输入是三个相邻的slice(类似像RGB通道一样处理)。因此,ResNet50的整个体系结构保持不变。多通道2D输入可以使2D网络在早期融合片间上下文。
GCN模块通过将一个大的K×K卷积核分解成两个一维核来模拟它,包含不同顺序排列的一维核的两个分支通过求和进行合并。每个Refine包含两个3×3的卷积,中间有一个Relu激活。每个Refine的输入也被添加到它的输出中,以形成一个残差连接。
Transferring the learned 2D net to 3D AH-Net
在二维MC-GCN网络收敛后,提取其编码器的参数,并将其传输到AH-net的相应编码器层。2D MC-GCN的译码器部分被丢弃,而解码部分则被丢弃。我们为AH-net设计了一个新的解码器,它由多层密集连接的块组成,然后是金字塔池化模块。
卷积张量从二维到三维的转换如下图所示,它的目的是在AH-net的编码器部分对3D卷逐片进行二维卷积,并且是将三个相邻的切片进行卷积。
然而,在三维空间中,Conv 1×1×1/(2,2,2)在z维上的每一步都跳过一个slice, 当图像只有少量的slice沿z 轴时,特别是用于检测任务的,这将丢失重要的信息。因此,作者使用Conv 1×1×1/(2,2,1)和MaxPool 1×1×2/(1,1,2)降采样来提取ResNet块之间的三维特征映射。
如上图中AH-Net ResNet Block所示,伴随着转换后的编码器,增加了三维解码器子网络,以利用三维图像上下文。在译码器中,具有Conv 1×1×1、Conv 3×3×1和Conv 1×1×3的卷积块。首先将特征传递到XY瓶颈块(不知道作者为什么称之为瓶颈块),Conv 3×3×1在两层Conv 1×1×1中间。然后输出到另一个中间为1×1×3的瓶颈块,并在转发到下一个块之前与其自身进行求和。这种卷积块将三维卷积分解为2D卷积和1D卷积。它首先使用2D卷积接收来自前几层的输入,保留了详细的2D特征,COV 1×1×3主要融合片内特征,使z维输出保持一致。
在预先训练的编码器和随机初始化的解码器上使用相同的学习速率来训练ah-net将使网络很难优化。为了训练3d AH-Net,首先锁定所有传输过来的参数,在优化过程中,只对译码器参数进行微调。然后,所有参数都可以一起微调。文中使用初始学习率0.0005对2D MCGCN进行微调。然后,在2D网络传输到3D后,将学习速率提高到0.001,对AH-Net进行微调,作者指出3D网络需要更大的学习速率才能在合理的时间内收敛。
文中Loss函数借鉴了Focal Loss,这个Loss函数与上面的金字塔池化一样,都是来自于何恺明大神的工作,他提出的这个损失函数,思路是希望那些hard examples对损失的贡献变大,使网络更倾向于从这些样本上学习。