TransFusionNet：JetsonTX2下肝肿瘤和血管分割的语义和空间特征融合框架

TransFusionNet：JetsonTX2下肝肿瘤和血管分割的语义和空间特征融合框架

IEEE JOURNAL OF BIOMEDICAL AND HEALTH INFORMATICS, VOL. 27, NO. 3, MARCH 2023

背景

CT图像中肝脏肿瘤和血管的分割和重建可以为医生的术前计划和手术干预提供便利。介绍了一个TransFusionNet框架，该框架由语义特征提取模块、局部空间特征提取模块，边缘特征提取模块和多尺度特征融合模块组成，以实现肝脏肿瘤和血管的细粒度分割。此外，我们将迁移学习方法应用于使用公共数据集进行预训练，然后对模型进行微调，以进一步提高拟合效果。此外，我们提出了一种智能量化方案来压缩模型权重，并在JetsonTX2上实现了高性能推理。TransFusionNet框架在血管分割任务中实现了0.854的平均IoU，在肝肿瘤分割任务中获得了0.927的平均IoU。在分析量化推理的计算性能时，我们的量化模型在NVIDIA RTX3090的Node上实现了4TFLOP，在JetsonTX2上实现了132GFLOP。这种前所未有的分割效果在一定程度上解决了自动分割的准确性和性能瓶颈。
需要解决的问题：
（1）如何设计一种准确、快速的肝肿瘤血管自动分割和三维重构方法？
（2） 如何设计一个能够学习空间语义融合特征的分割框架，以提高肿瘤和血管细节的分割精度？
（3） 如何设计高效的模型量化方法，实现肿瘤和血管的高性能模型推断和3D重建？
（4） 如何优化模型的计算和存储开销，以构建轻量级模型并将分割模型部署到JetsonTX2设备？

贡献

1）我们提出了TransFusionNet框架，该框架结合了CT图像的空间、语义和边缘特征，实现了对肝脏肿瘤和肝内动脉血管系统的精确精细分割。
2） 我们提出了一种基于强化学习的智能量化方案来压缩模型的权重，使模型在JetsonTX2和NVIDIA RTX3090 GPU的节点上都达到了最佳的推理性能。
3） 通过仔细量化，我们的模型使用NVIDIA RTX3090 GPU和JetsonTX2设备在Node上实现了高性能的肝血管肿瘤推断和三维重建。

实验

数据集：LITS（肝和肝肿瘤分割，https://competitions.codalab.org/competitions/17094)数据集包含130例肿瘤、转移瘤和囊肿。这些CT扫描具有较大的空间分辨率和视场（FOV）差异[4]。
3Dircadb（用于算法数据库比较的3D图像重建，https://www.ircad.fr/research/3d-ircadb-01/)是一个公共数据集，可用于训练和测试肝血管分割方法，包括20名不同图像分辨率、血管结构、强度分布和肝血管对比CT增强的患者[29]。
同时，我们收集了18例患者的CT增强图像，并构建了肝肿瘤血管（LTBV）数据集。为了清楚地区分肝脏中的血管并减轻标记负担，我们只保留了18名患者的动脉期图像。由于可训练样本的稀缺性，我们只分为训练数据集和测试数据集。我们使用这两个数据集进行模型预训练，并使用我们的私人数据集对肝动脉和肿瘤分割任务的模型进行微调训练。以上三个数据集按照8:2的比例分为训练集和测试集。

消融实验

方法

Transformer-Based Semantic Feature Extraction Module（transformer语义特征提取模块）

有效解决了传统深度CNN网络中感知场缺陷导致的信息缺失问题
可以学习全局特征表示的编码器，该编码器由基于特征提取主干的特征嵌入模块和基于Transformer感测图像的语义相关信息表示的特征提取模块组成[31]。该模块采用了一种全新的特征提取思想，通过对图片特征进行语义表示，并学习语义特征的全局表示。

Local Spatial Feature Extraction Module（局部空间特征提取模块）

基于Transformer的特征提取模块是一个非常强大的语义信息特征提取模块，因为Transformer特征提取模块在学习语义相关特征方面具有优势。然而，在许多方面，Transformer并不是传统卷积运算的有效替代品。对于提取一些图像中更细微的特征，如感兴趣区域的特征和微小血管的特征，CNN是一个完美的解决方案。我们设计了一种基于多层SEBottleNet堆叠的局部残差网络编码器。我们在BottleNet中引入了挤压和激励（SE）[51]，以增强特征图通道之间的相互依赖性。
确保每个模块都有不同的特征提取任务。我们在模块中间引入了挤压和激励模块，以更好地学习特征图通道维度的重要性，从而使SEBottleNet在特征提取过程中具有更强的学习重点。通过SEBottleNet和maxpool的连续叠加，编码器可以连续提取输入CT图像的局部特征表示。同时，由于每个SEBottleNet都设置了剩余连接，这使编码器能够有效地缓解网络深化带来的降级问题。

Edge Extraction Module（边缘提取模块）

由于肝动脉血管非常小，因此进一步细化血管和肝脏的分割是一项具有挑战性的任务。为了让模型能够学习更详细的空间特征，我们引入了EEM，它专门用于学习感兴趣的血管和肿瘤区域的边缘特征，并将边缘特征融合到分割网络中。

EEM以特征提取层的特征图和Canny算法[18]提取的CT边缘图（图3（b））为输入，预测边缘结果e∈RH×W。该模块预测边缘信息，并将预测的特征图组合到分割网络中。为了完成这项任务，我们处理分割注释以获得边缘注释er（图第3（d）段），它可以用作该模块的监督条件。

在本模块中，我们使用了门控激励卷积（GEC）层。GEC是EEM中最重要的单元，它可以过滤掉一些不相关的信息，集中精力提取图像的边缘特征。GEC应用于EEM和特征提取模块之间。它使用门控机制来停用其自身的激活，这些激活被提取模块[52]中包含的高级信息认为不相关。同时，我们在门控激活层中引入了一个激励模块，以学习不同特征图的重要性。

从理论上讲，GEC可以简单地看作是对特征图的空间维度和通道维度的关注的集合。通过GEC运算，注意力图αi选择性地保留了边缘语义特征。我们取消了对特征提取器的浅层特征图的GEC操作，因为馈送到卷积层的图像主要学习一般的低级特征，同时，输出的特征图保留了丰富的边缘信息。随着网络的深入，特征图将保留高级特征。在理论上，使用GEC运算可以有效地对高级特征的有用边缘信息进行加权。

Canny算子可以有效地滤除图像中不相关的特征，得到如图所示的Canny图像。第3（b）段。我们认为它适用于医学图像分割。因此，我们首先将canny图像和最后一个GEC模块输出en连接起来。然后，我们将它们与两个特征提取器的输出特征图一起提供给融合模块。同时，边缘提取模块以边缘损失为损失函数，以边缘标签为监督，对预测边缘图进行优化。

Multi-Scale Feature Fusing Module（多尺度特征融合模块）

多尺度特征融合解码模块，以对三个模块学习到的语义特征进行采样。该模块将三个模块提取的特征图作为输入，并输出预测的类别分布图

主要融合三个特征提取模块的特征图。参考空间金字塔池（SPP）来设计模块。首先，该模块使用C1×1和C3×3卷积分别从语义特征图和空间特征图的拼接结果中提取特征。接下来，我们将其输入到池化层，并融合边缘特征图。通过上述操作，获得了三个不同感受野的特征图。最后，我们对这三个特征图进行采样并连接，以输出融合的特征图。从理论上讲，这些模块输出的特征图可以保留丰富的空间特征、语义相关特征和边缘特征。在连续融合不同尺度的低级特征图后，解码器可以从粗到细地学习图像的语义信息。

Multi Task Training Strategy（多任务训练策略）

分割任务、边界任务、
优化过程中，特征提取模块和边缘提取模块的参数将基于损失进行优化。接下来，我们将边缘提取模块输出的特征图输入到融合模块中，以预测分割结果。因此，边缘提取模块学习到的先验知识保留在y中。同时，在优化过程中，分割损失会更加关注边缘特征

Applying Transfer Learning to TransFusionNet（迁移学习在TransFusionNet上的应用）

由于癌症筛查后增强CT图像的稀缺性以及肿瘤和血管的标记困难，我们获得了18例患者的CT图像。数据过少将不可避免地影响模型的性能，并加深过度拟合问题。为此，我们引入了一种迁移学习策略，该策略不需要具有确切代表性的训练数据，并且能够利用数据集之间的相似性，在模型的训练阶段捕获特定的先验知识，以构建新的分割模型。

首先使用公共数据集LITS和3Dircadb对模型进行预训练，分别获得肝脏肿瘤分割模型和肝脏血管分割模型。然后，我们使用我们的肝脏肿瘤数据和肝脏血管数据来重新训练通过预训练获得的模型。当我们需要对肝脏肿瘤和CT图像的血管进行分割时，我们只需要输入一张CT图像，模型就会分别对CT图像中的肿瘤和血管部分进行分割。

Quantification and Fine-Tuning of Inference Models（推理模型的量化与微调）

提出了一种基于硬件感知自动量化（HAQ）[55]的模型量化方案，以压缩框架的CNN和密集层，并优化计算和存储开销

损失函数

分割损失

边界损失

形状正则化损失，y是预测的分割图，e是预测的边缘图。在模型训练开始时，由于边缘提取模块是随机初始化的，无法准确预测e，（10）不起任何作用。因此，我们引入了一种动态调整策略，在100个历元之前将λ4设置为0，在100历元之后将λ4≥0

Thinking

双编码器分支，语义特征编码器在中间层加入Transformer，SEBottleNet分支提取局部特征，融合模块，边界特征提取模块，用标签的canny边缘图作为监督