CellViT实战指南：基于视觉Transformer的细胞分割与分类入门

1次阅读

共计 2307 个字符，预计需要花费 6 分钟才能阅读完成。

在生物医学图像分析中，细胞分割与分类是许多下游任务的基础。传统基于 CNN 的方法存在以下问题：

边缘模糊：卷积核的局部感受野难以建模长距离依赖关系，导致细胞边界分割不清晰
小目标漏检：下采样操作会丢失小尺寸细胞的细节信息
形变敏感：固定几何结构的卷积核难以适应细胞形态的多样性

视觉 Transformer(ViT) 通过自注意力机制解决了 CNN 的固有缺陷：

特性	CNN	ViT
感受野	局部	全局
位置编码	隐式 (通过卷积)	显式 (位置嵌入)
计算效率	高 (局部计算)	低 (序列长度平方)

CellViT 的创新点在于：

混合下采样结构：在 patch embedding 前保留 CNN 的局部特征提取能力
多尺度注意力：在不同层级建立远程依赖关系
类别感知损失：针对细胞分类任务优化损失函数

# 基础依赖
pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 
pip install albumentations==1.2.1 pytorch-lightning==1.7.7

import albumentations as A

train_transform = A.Compose([A.RandomRotate90(p=0.5),
    A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1, p=0.5),
    A.GaussianBlur(blur_limit=(3, 7), p=0.2),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.5),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])

import torch
import torch.nn as nn

class CellViT(nn.Module):
    def __init__(self, num_classes=3):
        super().__init__()
        # Patch embedding 层
        self.patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16)

        # Transformer 编码器
        encoder_layer = nn.TransformerEncoderLayer(d_model=768, nhead=12, dim_feedforward=3072)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=12)

        # 分割头
        self.seg_head = nn.Sequential(nn.ConvTranspose2d(768, 256, kernel_size=4, stride=4),
            nn.Conv2d(256, num_classes, kernel_size=1)
        )

    def forward(self, x):
        # 输入形状: [B, 3, 256, 256]
        patches = self.patch_embed(x)  # [B, 768, 16, 16]
        patches = patches.flatten(2).permute(2, 0, 1)  # [256, B, 768]
        features = self.transformer(patches)
        features = features.permute(1, 2, 0).view(-1, 768, 16, 16)
        return self.seg_head(features)

# 加权 Dice Loss 实现
class WeightedDiceLoss(nn.Module):
    def __init__(self, weights=[1.0, 2.0, 3.0]):
        super().__init__()
        self.weights = torch.tensor(weights).cuda()

    def forward(self, pred, target):
        smooth = 1.0
        pred = pred.softmax(dim=1)
        intersect = (pred * target).sum(dim=(2,3))
        denominator = (pred + target).sum(dim=(2,3))
        dice = (2. * intersect + smooth) / (denominator + smooth)
        loss = 1 - (dice * self.weights).mean()
        return loss

# 梯度累积（batch_size= 4 时等效于 bs=16）optimizer.zero_grad()
for i, (x, y) in enumerate(train_loader):
    pred = model(x)
    loss = criterion(pred, y)
    loss = loss / 4  # 梯度累积步数
    loss.backward()

    if (i+1) % 4 == 0:
        optimizer.step()
        optimizer.zero_grad()

在 MoNuSeg 测试集上的表现：