【医学分割】TransUnet

概述

transunet是transformer与unet的强强联合,transformer将CNN的特征图处理成sequence,借助self-attention操作捕捉全局的信息,将这部分信息上采样之后与高分辨率特征图融合,有效的提高了分割任务的效果,实现精准定位。
为什么要强强联合?
因为unet是优秀的分割网络,但是由于卷积操作天然的一些问题,比如平移不变性和捕捉长期依赖能力的不足等等,而transformer能够很好地解决这两个问题,但是它的处理缺乏一些细粒度信息,导致定位不够精确,所以两者的结合可能能够产生更好的效果。

细节

网络结构

首先将输入图片经过resnet50进行特征提取,其中的三个stage的输出将保留,用于之后的skip-connection。接着对于resnet50输出的特征图,将其进行序列化,送入transformer中进行序列预测,输出一个序列。然后将这个序列合并、reshape成一个新的特征图。接下来的过程和unet一样,就是不断上采样并和之前的resnet50的输出做skip-connection,最后还原回全分辨率的结果。
在这里插入图片描述

transformer中的部分细节

序列化操作
一个直观的序列化操作就是,对于一个HxWxC的图片,将每个像素作为一个vector,那么就会得到长度为一个HxW的sequence,但是这样的话,长度就有点大了。因此,一个新的做法就是,对于一个HxWxC的图片,选取N个patch或者说是N个小图,得到长度为N的sequence。假设patch的长度是p,那么 N = H ∗ W p ∗ p N=\frac{H*W}{p*p} N=ppHW,每个vector中元素的个数就是 p ∗ p ∗ C p*p*C ppC
transformer encoder的输出
然后这里用的是transformer的encoder部分,所以输入和输出的一样长度的,因此输出是N个 p ∗ p ∗ C p*p*C ppC的vector,我们将输出拼接起来能够得到一个 H ∗ W p ∗ p ∗ ( p ∗ p ∗ C ) \frac{H*W}{p*p}*(p*p*C) ppHW(ppC)的矩阵,我们将它reshape成 H p ∗ W p ∗ ( p ∗ p ∗ C ) \frac{H}{p}*\frac{W}{p}*(p*p*C) pHpW(ppC)的tensor,接着使用1x1的卷积对通道数进行变换,那么就相当于得到了encoder最后一层的输出。

CNN-Transformer混合的结构

作者提了两种思路,一种是Transformer直接对原图做处理,然后将encoder的结果直接上采样到全分辨率,但是效果并不出色,可能的原因是 H p ∗ W p \frac{H}{p}*\frac{W}{p} pHpW的尺寸相对于 H ∗ W H*W HW是有点小了,那么直接上采样的话,缺少很多的细粒度信息。
因此作者提出了第二种思路,就是论文中讲的这种,一方面Transformer不直接对原图做处理,而是对特征图做处理,另一方面将一些细粒度信息通过skip-connection融合过来。

级联上采样

在对transformer encoder部分的输出做处理之后,将其不断上采样的过程中,并不是常规的上采样步骤而是级联的上采样,每个上采样块是常规的上采样、3x3卷积还有relu。

### TransUNet工作原理 TransUNet融合了Transformer和传统卷积神经网络的优点,在医学图像分割领域取得了显著进展。该架构通过引入全局自注意力机制来增强特征提取能力,同时保留局部细节信息。 #### 1. 结合CNN与Transformer的优势 为了克服单一模型的局限性,TransUNet采用了一种混合编码器设计。具体来说: - 初始阶段利用ResNet作为基础骨干网,负责捕捉低层次的空间特征[^1]。 ```python import torch.nn as nn class ResidualBlock(nn.Module): def __init__(self, in_channels, out_channels): super(ResidualBlock, self).__init__() self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) self.bn1 = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out += residual out = self.relu(out) return out ``` #### 2. Transformer层的设计 随后加入多头自注意力模块(Multi-head Self-Attention),使得模型能够关注到更广泛的上下文关系,从而改善远距离依赖的学习效果[^2]。 - 使用线性投影将输入映射至较低维度空间; - 计算Query、Key 和 Value矩阵之间的相似度得分; - 应用Softmax函数并加权求和得到最终输出表示; ```python import math from typing import Optional def scaled_dot_product_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: d_k = query.size(-1) scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) p_attn = F.softmax(scores, dim=-1) output = torch.matmul(p_attn, value) return output ``` #### 3. 解码路径优化 解码部分沿用了经典的跳跃连接策略,即从编码端获取对应尺度下的语义信息并与当前层级相拼接,以此恢复高分辨率细节[^3]。 ```python class DecoderBlock(nn.Module): def __init__(self, in_channels, mid_channels, out_channels): super().__init__() self.upconv = nn.Sequential( nn.Upsample(scale_factor=2), nn.Conv2d(in_channels, mid_channels, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(mid_channels), nn.ReLU() ) self.double_conv = DoubleConv(mid_channels * 2, out_channels) def forward(self, encoder_features, decoder_features): upsampled_decoder_features = self.upconv(decoder_features) concatenated_features = torch.cat([encoder_features, upsampled_decoder_features], dim=1) return self.double_conv(concatenated_features) class DoubleConv(nn.Module): """(convolution => [BN] => ReLU) * 2""" def __init__(self, in_channels, out_channels): super().__init__() self.double_conv = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU(), nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), nn.BatchNorm2d(out_channels), nn.ReLU() ) def forward(self, x): return self.double_conv(x) ```
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值