将 Vision Transformer 用于医学图像的语义分割

关于ViT的关键点如下：

ViT架构基于将图像表示为一组补丁。图像补丁是图像的非重叠块。每个块最初都有一个由该块中的图像像素形成的嵌入向量。
Transformer编码器是ViT的主要部分，它根据它们的类别归属来训练补丁之间的相似度。它包含一系列线性、归一化和激活层。
在大型数据集（例如ImageNet21K）上预训练的ViT模型可以用于在自定义数据集上进行迁移学习，微调后的模型表现良好。

关于U-Net的关键点如下：

U-Net由两部分组成：编码器和解码器。编码器包含一系列用于特征提取和图像缩减的块。解码器对称于编码器，重构图像分辨率。
在CNN中，U-Net是医学图像语义分割的最佳架构之一。

我在我的分割系统中使用了Hugging Face的Swin Transformer V2作为编码器。Swin Transformer（分层视觉Transformer，使用偏移窗口）包含4个阶段的编码器处理嵌入补丁。最初，补丁大小为4x4像素。在每个编码器阶段，通过合并来自前一个阶段较小补丁的嵌入，补丁分辨率会增加两倍。这意味着图像的空间分辨率，以补丁表示，每个后续阶段会减少两倍。下图（来自Hugging Face文档）显示了Swin Transformer的高级架构：

请注意，带有下采样编码器块序列类似于我之前文章中讨论的U-Net编码器的高级架构。还请注意，用于分类的ViT编码器在任何阶段都使用16x16的补丁（参见上图）。

已经训练了用于分割的Swin Transformer的几个模型，包括一个在ImageNet21K数据集上训练的大型模型（~ 1400万张图像）。完整的分割流水线由编码器和解码器组成。使用Hugging Face的Swin Transformer编码器进行以下自定义数据集的微调。换句话说，我使用预训练的Swin Transformer大型模型作为编码器，并实现和训练我的自定义解码器，以构建用于我的数据集的语义分割的完整系统。

从Hugging Face加载的Swin Transformer V2：深入了解

让我们使用以下代码块查看来自Hugging Face的Swin Transformer V2模型：

安装

在这里插入图片描述

导入

在这里插入图片描述

谷歌云硬盘挂载（用于Google Colab）

在这里插入图片描述

CUDA设备配置

device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

大预训练模型的加载（在ImageNet21K上训练）

from transformers import AutoImageProcessor, Swinv2Model

image_processor = AutoImageProcessor.from_pretrained("microsoft/swinv2-large-patch4-window12-192-22k")
model = Swinv2Model.from_pretrained("microsoft/swinv2-large-patch4-window12-192-22k")

image_processor 定义了一组应用于输入图像的变换，这些图像最初以PIL图像的形式存在：

在这里插入图片描述

将输入的PIL图像转换为torch张量，将其调整为图像分辨率192x192，并进行归一化。

模型摘要

summary(model=model, input_size=(1, 3, 192, 192), col_names=['input_size', 'output_size', 'num_params', 'trainable'])

这是一个包含超过1.95亿个参数的大型模型。

调用

model.eval()

以查看模型包含的所有层。

Swin Transformer V2模型的部分如下

Patch-embeddings 层，将输入图像的分辨率192x192划分为2034=48*48个4x4大小的patch。对于每个patch，形成一个长度为192的线性投影向量。
4个编码器阶段。在每个阶段，进行Multi Head Self Attention的训练。每个阶段中像素大小的补丁会两倍增加（通过补丁嵌入合并）。每个阶段中，以补丁表示的图像分辨率会减少两倍。
对编码器输出进行归一化，以生成来自编码器的last_hidden_state。
对last_hidden_state张量进行平均池化，以生成包含类别嵌入的pooler_output向量。

加载任何图像并通过image_processor对其进行预处理：

在这里插入图片描述

将输入发送到Swin Transformer V2模型。下图说明了按顺序调用Swin Transformer V2的各个部分（如左图所示）等效于将整个模型作为一个整体进行调用（如右图所示）：

查看每个编码器阶段的输出：

在这里插入图片描述

我们看到以下形状：

im0 -> torch.Size([1, 2304, 192]) -> 2304=48*48 — patch数，192 — patch-embeddings长度
im1 -> torch.Size([1, 576, 384]) -> 576=24*24 — patch数，384 — patch-embeddings长度
im2 -> torch.Size([1, 144, 768]) -> 144=12*12 — patch数，768 — patch-embeddings长度
im3 -> torch.Size([1, 36, 1536]) -> 36=6*6 — patch数，1536 — patch-embeddings长度
im4 -> torch.Size([1, 36, 1536]) -> 36=6*6 — patch数，1536 — patch-embeddings长度

这些来自Swin Transformer V2预训练大型模型的输出将成为我的解码器模型的输入。我训练我的解码器以获得大脑MRI异常区域的分割掩模。下图显示了此解码器的高级架构：

请注意，在上图流程图的最后一个块中，“图像调整大小为256x256”是流程图中常规解码器流程的自定义元素：我将大脑MRI的图像分辨率设置为256x256，并将其应用于分割掩模图像。

基于Swin Transformer V2的大脑MRI语义分割系统的实现

让我们回到大脑MRI的语义分割任务。我使用了来自Kaggle的大脑MRI数据集。该数据集包含110位患者的数据：每位患者都有一组带有脑部切片的MRI图像，以及相应的带有异常区域掩码图像的图像。下图显示了“脑部切片图像 + 带有掩码图像”的示例对：

在数据集中，每个人的“脑切片图像+掩模图像”对的数量从20对到88对不等。整个集合包含 3935 对：2556 对具有零掩模，1379 对具有针对异常区域的非零掩模。

我使用PyTorch对基于Swin Transformer V2阶段输出（im0、im1、im2、im3、im4 — 见图1）的自定义解码器模型进行实现和训练。下面的代码显示了对一个图像进行预处理，该图像最初以PIL图像的形式存在，然后转换为来自预训练的Swin Transformer V2模型阶段的im0、im1、im2、im3、im4张量。下面代码中的变量model是加载的在ImageNet21K上预训练的Swin Transformer V2模型（请参阅上一节的代码块）：

img = <load PIL Image>
img = image_processor(images=img, return_tensors="pt")

x = model.embeddings(**img)
input_dimensions=x[1]
im0 = x[0].detach().squeeze()

x = model.encoder.layers[0](x[0], input_dimensions=input_dimensions)
im1 = x[0].detach().squeeze()

x = model.encoder.layers[1](x[0], input_dimensions=(input_dimensions[0]//2, input_dimensions[1]//2))
im2 = x[0].detach().squeeze()

x = model.encoder.layers[2](x[0], input_dimensions=(input_dimensions[0]//4, input_dimensions[1]//4))
im3 = x[0].detach().squeeze()

x = model.encoder.layers[3](x[0], input_dimensions=(input_dimensions[0]//8, input_dimensions[1]//8))
x = model.layernorm(x[0])
im4 = x.detach().squeeze()

注意：我使用squeeze()来删除单个图像的批处理维度，因为我假设它将被发送到torch-DataLoader，后者会向图像批次添加批处理维度。

仅对掩模图像应用转换为torch张量和调整大小。来自DataLoader的5个输入张量被发送到以下模型：

class Up_Linear(nn.Module):
    def __init__(self, in_ch, size, coef=1):    
        super(Up_Linear, self).__init__()        
        self.shuffle = nn.PixelShuffle(upscale_factor=2)      
          
        n_ch = int(coef * in_ch)                
        
        self.ln = nn.Sequential(      
            nn.Linear(in_ch * 2, n_ch),            
            nn.ReLU(inplace=True),            
            nn.Linear(n_ch, in_ch * 2),            
            nn.ReLU(inplace=True),        
        )                
        
        self.size = size    
    def forward(self, x1, x2):   
        x = torch.cat((x1, x2), 2)        
        x = self.ln(x)        
        x = x.permute(0, 2, 1)        
        x = torch.reshape(x, (x.shape[0], x.shape[1], self.size, self.size))        
        x = self.shuffle(x)        
        x = torch.reshape(x, (x.shape[0], x.shape[1], self.size*self.size*4))        
        x = x.permute(0, 2, 1)        
        return x

class MRI_Seg(nn.Module):
    def __init__(self):    
        super(MRI_Seg, self).__init__()        
        
        self.ups3 = Up_Linear(1536, 6, 1)        
        self.ups2 = Up_Linear(768, 12, 1)        
        self.ups1 = Up_Linear(384, 24, 2)        
        self.ups0 = Up_Linear(192, 48, 3)      
          
        self.shuffle = nn.PixelShuffle(upscale_factor=2)    
            
        self.out = nn.Sequential(      
            nn.Conv2d(24, 1, kernel_size=1, stride=1),            
            nn.Sigmoid()        
        )    
        
    def forward(self, x0, x1, x2, x3, x4):
        x = self.ups3(x4, x3)        
        x = self.ups2(x, x2)        
        x = self.ups1(x, x1)        
        x = self.ups0(x, x0)      
          
        x = x.permute(0, 2, 1)        
        x = torch.reshape(x, (x.shape[0], x.shape[1], 96, 96))        
        x = self.shuffle(x)        
        x = transforms.Resize((256, 256))(x)       
         
        x = self.out(x)        
        return x  


net = MRI_Seg().to(device)

模型摘要

summary(model=net, input_size=[(1, 2304, 192), (1, 576, 384), (1, 144, 768), (1, 36, 1536), (1, 36, 1536)], col_names=['input_size', 'output_size', 'num_params', 'trainable'])

该模型包含超过1300万个可训练参数（类似于U-Net模型）。

我使用二元交叉熵损失函数来训练我的模型，以使生成的掩码更接近标签（掩码）图像。我使用Adam优化器和学习率0.0001。我同时使用IoU（交并比）和Dice指标作为质量度量：IoU = 1和Dice = 1表示理想质量。请注意，对于包括图片在内的所有结果，我使用经过训练模型生成的分割掩码，并应用阈值：如果掩码像素值小于0.5，则将掩码像素值设置为0，否则将掩码像素值设置为1。

下图显示了对U-Net架构（它们是在此处获得和呈现的）和MRI_Seg模型（上面）的结果进行比较。对于两种模型，我选择了在测试集上显示最佳结果的检查点。对于对U-Net和基于Transformer的模型的性能评估，我使用了来自训练集的550个随机选择的条目和所有394个包含在测试集中的条目（使用相同的训练和测试集）：

正如我们所见，基于Transformer的分割模型的性能明显低于U-Net模型的性能。我尝试过改变我的解码器架构和可训练参数的数量，但并未改善性能。

我对MRI和掩码图像应用了水平翻转（改变图像的左右侧），并创建了一个图像数量加倍的数据集。现在，我正在使用这个增强的数据集来训练基于Transformer的模型。下图显示了U-Net架构（在此处获取和呈现）和MRI_Seg模型（上面）在增强数据集上结果的比较。对于两种模型，我选择了在测试集上表现最佳的检查点。对于基于U-Net和Transformer的模型的性能评估，我使用了从训练集中随机选择的550个条目以及测试集中包含的所有787个测试条目（两种模型都使用相同的训练和测试集）：

我们可以看到，基于增强数据集训练的基于Transformer的分割模型的性能显着提高，并且接近于U-Net模型的性能。

下面的图片显示了基于Transformer的训练模型在测试图像上的工作结果。

结论

Swin Transformer V2的高级架构允许使用U-Net模型的概念来实现自定义解码器的语义分割系统。
自定义数据集的大小应该足够大（约为10K）。在这种情况下，基于Transformer的语义分割系统的性能良好。
基于Transformer的语义分割系统能够获得接近于最佳分割模型（如U-Net）的性能。

· END ·

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述