[论文阅读] Polyp-PVT: Polyp Segmentation with Pyramid Vision Transformers

论文地址:https://arxiv.org/abs/2108.06932
代码:https://github.com/DengPingFan/Polyp-PVT
发表于:CAAI AIR 23

Abstract

大多数息肉分割方法使用CNN作为其骨干,导致在编码器和解码器之间交换信息时出现两个关键问题:1)考虑到不同级别特征之间的贡献差异;2)设计一个有效的机制来融合这些特征。与现有的基于CNN的方法不同,我们采用了一个Transformer编码器,它可以学习到更强大和稳健的表征。此外,考虑到图像采集的影响和息肉难以捉摸的特性,我们引入了三个新的模块,包括级联融合模块(CFM)、伪装识别模块(CIM)和相似度聚合模块(SAM)。其中,CFM用于从高层特征中收集息肉的语义和位置信息,而CIM则用于捕捉伪装在低层特征中的息肉信息。在SAM的帮助下,我们将具有高层语义位置信息的息肉区域的像素特征扩展到整个息肉区域,从而有效地融合了跨层特征。所提出的模型被命名为Polyp-PVT,它有效地抑制了特征中的噪音,并显著提高了它们的表达能力。在五个广泛采用的数据集上进行的广泛实验表明,与现有的方法相比,所提出的模型对各种具有挑战性的情况(如外观变化、小物体)更加稳健,并实现了新的最先进的性能。

I. Method

摘要可能写的已经比较直白了,本文的最大特点就是用上了Transformer Encoder,其他的大创新…可能没有。
在这里插入图片描述
Encoder换成了PVT(v2),而Decoder的话,这里放弃了U-Net那种逐级上采样的思路,而是取出Encoder不同层的特征各自进行处理,然后最后给揉起来,不过这种思路在某个二分割任务重也比较常见。那么接下来我们就简单提一下本文所提出的三个模块。

II. Cascaded Fusion Module(CFM)

在这里插入图片描述
有点拿更高级特征为低级特征做attention的意思,叠buff叠的有点复杂,这里就不细讲了,可以看图。这里需要注意的是做了一个事实上的中间decoder输出来构造side output supervision。

III. Camouflage Identification Module(CIM)

在这里插入图片描述
串联通道注意力和空间注意力…这个东西很早就有人专门研究过了[1],而且应用也极其广泛,只能说文中把这个搞成像自己提出的东西一样还另外取了个名字实在是不妥(好歹改一改)。不过也再次证明通道注意力和空间注意力这俩东西确实是刷点的好玩意。

IV. Similarity Aggregation Module(SAM)

在这里插入图片描述
看起来很复杂。回到最开始的网络整体结构图的话,SAM接受了两个输入,一个是CFM模块输出的高级特征 T 1 T_1 T1,一个是SIM模块输出的低级特征 T 2 T_2 T2,那么从输入就可以看出来这个东西做的还是多级特征融合。而多级特征融合,解决的其实就是怎么加权然后揉起来的问题,这里采用的是Transformer中的Self Attention思想,那么就成了怎么构造Q、K、V的问题。

具体操作上,对于高级特征 T 1 T_1 T1,首先走两个1×1卷积进行降维,得到了 Q Q Q K K K。对于低级特征 T 2 T_2 T2,也是类似的处理方式,得到 V V V K K K先与 V V V相乘,然后再与 Q Q Q相乘,至此便完成了Self Attention的过程。这里比较特别的是,对于融合后的特征,使用了一个GCN来进一步提取相关性。

总的来说可以认为是Non-Local、Self Attention、GCN的结合。

V. Loss

从第II节中可以看到,网络实际上是有个“辅助输出”的,但是这个东西并不会作为最终的预测结果,而单纯只是用来对网络训练进行监督。具体来说,网络的损失函数如下: L = L main  + L aux  \mathcal{L}=\mathcal{L}_{\text {main }}+\mathcal{L}_{\text {aux }} L=Lmain +Laux  main的话就是对最终结果 P 2 P_2 P2的监督,而aux的话就是对中间结果 P 1 P_1 P1的输出,有: L main  = L I O U w ( P 2 , G ) + L B C E w ( P 2 , G ) \mathcal{L}_{\text {main }}=\mathcal{L}_{\mathrm{IOU}}^{w}\left(P_{2}, G\right)+\mathcal{L}_{\mathrm{BCE}}^{w}\left(P_{2}, G\right) Lmain =LIOUw(P2,G)+LBCEw(P2,G) L a u x = L I o U w ( P 1 , G ) + L B C E w ( P 1 , G ) \mathcal{L}_{\mathrm{aux}}=\mathcal{L}_{\mathrm{IoU}}^{w}\left(P_{1}, G\right)+\mathcal{L}_{\mathrm{BCE}}^{w}\left(P_{1}, G\right) Laux=LIoUw(P1,G)+LBCEw(P1,G) 其实用的是一样的。

Ref

[1] Woo, S., Park, J., Lee, J. Y., & Kweon, I. S. (2018). Cbam: Convolutional block attention module. In Proceedings of the European conference on computer vision (ECCV) (pp. 3-19).

### Polyp-Mamba 代码实现概述 Polyp-Mamba 是一种基于状态空间模型 (SSM) 的方法,专门用于医学图像处理中的息肉分割任务。该模型通过结合二元交叉熵损失 \(L_{BCE}\) 和加权 IoU 损失 \(L_{Iou}\),并通过自适应权重调节实现了高性能和快速收敛[^3]。 以下是 Polyp-Mamba 的核心组件及其对应的代码片段: #### 总体结构 Polyp-Mamba 结合了一个编码器-解码器架构,其中编码器负责特征提取,而解码器则通过跳跃连接恢复细节并生成最终的分割图[^4]。 ```python import torch.nn as nn import torch class Encoder(nn.Module): def __init__(self, channels=[64, 128, 256]): super(Encoder, self).__init__() self.patch_merging_layers = nn.ModuleList([ PatchMerging(channels[i]) for i in range(len(channels)-1) ]) def forward(self, x): features = [] for layer in self.patch_merging_layers: x = layer(x) features.append(x) return features class Decoder(nn.Module): def __init__(self, channels=[256, 128, 64], skip_channels=[256, 128, 64]): super(Decoder, self).__init__() self.patch_expanding_layers = nn.ModuleList([ PatchExpanding(skip_channels[i], channels[i]) for i in range(len(channels)) ]) def forward(self, encoder_features): x = None for idx, layer in enumerate(reversed(self.patch_expanding_layers)): if x is None: x = encoder_features[-idx-1] else: x = layer(torch.cat([x, encoder_features[-idx-1]], dim=1)) return x class FinalProjectionLayer(nn.Module): def __init__(self, input_channel, output_channel): super(FinalProjectionLayer, self).__init__() self.up_sample = nn.Upsample(scale_factor=4, mode='bilinear', align_corners=True) self.projection = nn.Conv2d(input_channel, output_channel, kernel_size=1) def forward(self, x): x = self.up_sample(x) x = self.projection(x) return x class PolypMambaModel(nn.Module): def __init__(self, enc_chs=[64, 128, 256], dec_chs=[256, 128, 64], n_classes=1): super(PolypMambaModel, self).__init__() self.encoder = Encoder(enc_chs) self.decoder = Decoder(dec_chs[::-1], enc_chs[:-1]) self.final_projection = FinalProjectionLayer(dec_chs[0], n_classes) def forward(self, x): encoded_features = self.encoder(x) decoded_output = self.decoder(encoded_features) final_output = self.final_projection(decoded_output) return final_output ``` #### 损失函数定义 为了优化 Polyp-Mamba 的性能,采用了组合损失函数的形式,即二元交叉熵损失和加权 IoU 损失的线性组合。 ```python def binary_cross_entropy_loss(pred, target): bce_loss_fn = nn.BCEWithLogitsLoss() return bce_loss_fn(pred, target) def weighted_iou_loss(pred, target): intersection = (pred * target).sum(dim=(1, 2, 3)) union = pred.sum(dim=(1, 2, 3)) + target.sum(dim=(1, 2, 3)) - intersection epsilon = 1e-7 loss = ((intersection + epsilon) / (union + epsilon)).mean() return 1 - loss def total_loss(pred, target, lambda1=1.0, lambda2=1.0): bce_loss = binary_cross_entropy_loss(pred, target) iou_loss = weighted_iou_loss(pred, target) return lambda1 * bce_loss + lambda2 * iou_loss ``` ### §相关问题§ 1. 如何调整超参数 \(\lambda_1\) 和 \(\lambda_2\) 来进一步提升 Polyp-Mamba 的性能? 2. 在实际部署中,如何加速 Polyp-Mamba 的推理过程而不显著降低精度? 3. 是否可以通过迁移学习将 Polyp-Mamba 扩展到其他医疗影像分割任务? 4. 如果训练数据不足,能否利用半监督或无监督技术增强 Polyp-Mamba 的泛化能力? 5. 对于大规模数据集,如何优化 Polyp-Mamba 的训练效率?
评论 13
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值