Mask2former源码解读

在这里插入图片描述

CVPR2022

整个网络可以分成四大部分:Backbone ; Pixel Decoder ; Transformer Decoder ; 最后的输出

接下来串一下整个源码部分

首先需要明确:

  1. input: image
  2. output: class 和 mask

1. train_net.py

执行主程序:首先解析命令行参数,args保存了从命令行传入的配置信息

接下来调用launch()函数来启动整个分布式训练流程,可以看到很明显得启动主函数main()

if __name__ == "__main__":
    args = default_argument_parser().parse_args()  # 解析命令行参数
    print("Command Line Args:", args)
    launch(
        main,  # 调用训练主函数
        args.num_gpus,  # GPU数量
        num_machines=args.num_machines,
        machine_rank=args.machine_rank,
        dist_url=args.dist_url,
        args=(args,),
    )

这是训练的主程序入口

def main(args):
    cfg = setup(args)  # 根据args来初始化配置对象cfg

    #(1)在评估模式下
    if args.eval_only:
        model = Trainer.build_model(cfg) # 根据配置对象 cfg 构建模型 model
        DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
            cfg.MODEL.WEIGHTS, resume=args.resume
        )# 不用管(关于恢复训练的)
        res = Trainer.test(cfg, model) # 训练好的model在测试集上的结果(记录评估结果)
        if cfg.TEST.AUG.ENABLED:
            res.update(Trainer.test_with_TTA(cfg, model))
        if comm.is_main_process():
            verify_results(cfg, res)
        return res

    #(2)在训练模式下
    trainer = Trainer(cfg)  # 根据配置信息cfg来创建一个训练器实例trainer
    trainer.resume_or_load(resume=args.resume) 
    
    # 返回训练结果(训练结果也就是训练好的model)
    return trainer.train()

整个流程来看的话肯定是先训练model 训练好之后再去利用训练好的model在测试集上测试 得到测试结果(也就是评估结果)

1.首先:是调用setup()函数

根据我们上一步解析的命令行参数args来初始化和配置模型训练所需的各项配置,返回一个配置对象cfg,之后我们读取配置文件得到各个部分所需要的配置信息都存储在这里。

def setup(args):
    
    # (1)初始化配置
    cfg = get_cfg()  # 初始化一个基础配置对象cfg
    
    # (2)扩展配置功能
    add_deeplab_config(cfg)  # 向配置对象cfg中添加与DeepLab模型相关的特定配置项
    add_maskformer2_config(cfg)  # 向配置对象cfg中添加MaskFormer v2模型的特定配置
    
    # (3)合并外部配置
    cfg.merge_from_file(args.config_file) # 从用户指定的配置文件中读取额外的配置信息
    cfg.merge_from_list(args.opts) # 用户通过命令行参数传递的额外配置
    
    # (4)冻结配置
    cfg.freeze()
    default_setup(cfg, args)

    setup_logger(output=cfg.OUTPUT_DIR, distributed_rank=comm.get_rank(), name="mask2former")
    
    return cfg
2. 然后

根据输入的参数 args 来决定执行模型的训练还是评估过程。

主要还是关注训练模式:
定义了Trainer类这个类继承了基类(DefaultTrainer类,分别负责根据配置信息cfg来创建模型的优化器和训练数据加载器等。)
在这里插入图片描述

2.maskformer_model.py

在这里插入图片描述
把整个系统的各个部分串起来
在这里插入图片描述
在这里插入图片描述

第一部分是 Backbone(swin.py文件)

1. 骨干网络做了啥??

传入经过预处理后的图片image( B ; 3 ; H ; W )========>经过一系列处理得到四个尺寸的特征图(res2 ; res3 ; res4 ;res5 )

在这里插入图片描述

2. 骨干网络是怎么来的??怎么构建的??

首先根据 cfg 配置参数来构建 backbone ,函数从配置文件中的MODEL.BACKBONE.NAME读取骨干网络的名称

在这里插入图片描述

然后 调用BACKBONE_REGISTRY.get(backbone_name),根据配置中指定的backbone_name查找对应的创建函数,然后通过这个函数和给定的配置(cfg)及输入形状(input_shape)来实例化一个具体的骨干网络对象 D2SwinTransformer

3.骨干网络是怎么工作的??为啥照片进去就得到特征图了??
  • 逐步特征提炼: 从原始图像到特征图的转化是一个逐步抽象的过程。每经过一层,图像的信息被转换成更高层次、更具语义性的特征表示。Swin Transformer通过局部窗口内的注意力机制和跨层的信息整合,逐步提炼出对物体类别、形状、纹理等关键属性敏感的特征图。这些特征图不再直接关联原始像素,而是代表了图像内容的高层次理解

D2SwinTransformer前向传播过程:

# 前向传播
def forward(self, x):
    # 断言确保输入x是一个四维张量,形状为(N, C, H, W)
    assert (
        x.dim() == 4
    ), f"SwinTransformer takes an input of shape (N, C, H, W). Got {x.shape} instead!"

    # (1)初始化空字典,存储backbone处理后特征图(就是O1;02;03;04)
    outputs = {}
    # (2)调用基类的forward方法处理输入x**核心
    y = super().forward(x)
    # (3)从各种分辨率的特征图中筛选出O1;O2;O3;O4来
    for k in y.keys():
        if k in self._out_features:
            outputs[k] = y[k]
    return outputs

SwinTransformer基类前向传播:


def forward(self, x):
    
    # (1)图像分割成小块 patch,然后每个小方块都会被转化成一个特征向量(让机器理解patch)
    x = self.patch_embed(x)

    # (2)为每个patch的特征向量添加位置编码(这样机器就知道每个patch在整个图片中的位置了)
    Wh, Ww = x.size(2), x.size(3) 
    if self.ape:
        absolute_pos_embed = F.interpolate(
            self.absolute_pos_embed, size=(Wh, Ww), mode="bicubic"
        )
        x = (x + absolute_pos_embed).flatten(2).transpose(1, 2)  # B Wh*Ww C
    else:
        x = x.flatten(2).transpose(1, 2)
        
    # 特征图 x 进行dropout处理,以防过拟合
    x = self.pos_drop(x)    

    # (3)每一层都会从之前的特征中提取更深层、更抽象的信息
    outs = {}  # 存放O1;O2;O3;O4
    # 遍历网络中的每一层(self.layers)
    for i in range(self.num_layers):
        layer = self.layers[i]
        x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
        # 如果当前层的索引i位于self.out_indices中,表示该层的输出需要被收集:
        if i in self.out_indices:
            norm_layer = getattr(self, f"norm{i}")
            x_out = norm_layer(x_out)	# 对输出进行归一化处理
            # 调整x_out的形状以匹配预期的输出格式
            out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
            outs["res{}".format(i + 2)] = out # 存储输出特征图

    # (4)# 返回所有指定输出层的特征图
    return outs 

图像分割与嵌入

PatchEmbed 将输入图像分割成多个不重叠的Patch,对每个Patch进行线性映射以获得高维嵌入向量,并根据配置进行归一化处理,最终输出特征图,为后续的Transformer结构提供合适的输入格式

# 前向传播过程
def forward(self, x):
    # 获取输入图像的尺寸
    _, _, H, W = x.size()
    # (1)如果图像的宽度或高度不能被patch_size整除,则进行填充
    if W % self.patch_size[1] != 0:
        # 在宽度方向上进行右填充
        x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
    if H % self.patch_size[0] != 0:
        # 在高度方向上进行下填充
        x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))

    # (2)使用卷积层将图像分割成Patch并映射到嵌入空间
    x = self.proj(x)  # 输出形状变为(Batch, embed_dim, new_height, new_width)

    # (3) 如果存在归一化层,则对输出进行归一化处理
    if self.norm is not None:
        # 首先将形状从(B, C, H, W)调整为(B, H*W, C),以便进行归一化
        Wh, Ww = x.size(2), x.size(3)
        x = x.flatten(2).transpose(1, 2)
        x = self.norm(x)
        # 再将形状调整回(B, C, H, W)以便于后续操作
        x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)

    return x

在上述前向传播过程中遍历各层进行特征提取(该层中的所有SwinTransformerBlock

首先需要一个函数 负责串联多个SwinTransformerBlock并进行下采样操作(如果配置)

假设:::

想象现在有一张大照片,我们想让机器仔细查看并理解这张照片的每个部分。但是,为了让机器更高效地工作,我们把它分成很多小格子。

首先要画一个计划图,告诉机器哪些格子是相邻的,也就是“注意力掩码attn_mask”。

  • 先根据原始图片的尺寸算出需要多少行和列的小格子,确保整个照片都被覆盖,在边缘补上空白。
  • 然后,开始在每个小格子上做标记。每个相邻或相隔特定距离的格子会得到不同的编号,这样机器人就知道哪些格子应该一起看,哪些应该分开考虑。
  • 编号相同的格子(代表同一区域内的小块)可以互相“看见”,而不同编号的格子则假装对方不存在(通过设置一个非常低的分数,比如-100)。

接下来,你有一组可以学习的机器(每个Swin Transformer Block就是一个这样的机器),它们会依次查看并研究这张照片上的每个小格子。

  • 为了不耗尽机器内存,有时候你会用一种叫做“检查点”的技术,让机器人只记住最重要的步骤,忘记一些中间过程。

如果需要,你还可以告诉机器,下次再看照片时,每个小格子要包含更多的信息,也就是说,让格子变大一些。这就是下采样处理——它帮助机器从更宏观的角度理解照片,不过这一步不是每次都需要做。

###前向传播
def forward(self, x, H, W):
    #(1) 构建注意力掩码(告诉机器哪些小块是相邻的)
    Hp, Wp = int(np.ceil(H / self.window_size)) * self.window_size, int(
        np.ceil(W / self.window_size)) * self.window_size  # 补齐后的分辨率
    # 初始化一个与补足分辨率匹配的全零掩码张量`img_mask`
    img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  
    # 分割掩码以匹配窗口大小和移位
    h_slices, w_slices = [(slice(0, -self.window_size), slice(-self.window_size, -self.shift_size),
                           slice(-self.shift_size, None))] * 2
    cnt = 0
    for h in h_slices:
        for w in w_slices:
            img_mask[:, h, w, :] = cnt  # 分配唯一ID给掩码的每个窗口
            cnt += 1
            
    # 将掩码按窗口大小划分并调整形状
    mask_windows = window_partition(img_mask, self.window_size)
    mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
    # 构建实际的注意力掩码`attn_mask`
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
    attn_mask = attn_mask.masked_fill(attn_mask != 0,float(-100.0)).masked_fill(attn_mask == 0,float(0.0))  # 不同窗口间的元素设为负无穷大(`-100.0`),表示不应关注;同一窗口内的元素设为0,表示可以正常参与注意力计算。
    
    # (2)遍历并执行每个SwinTransformerBlock
    for blk in self.blocks:  # 遍历该层中的所有`SwinTransformerBlock`,对每个块执行前向传播
        blk.H, blk.W = H, W  # 设置当前块的空间分辨率
        if self.use_checkpoint:
            x = checkpoint.checkpoint(blk, x, attn_mask)  # 使用检查点以节省内存
        else:
            x = blk(x, attn_mask)  # 前向传播

    #(3) 下采样处理
    if self.downsample is not None:
        x_down = self.downsample(x, H, W)  # 执行下采样
        Wh, Ww = (H + 1) // 2, (W + 1) // 2  # 下采样后的分辨率
        return x, H, W, x_down, Wh, Ww  # 返回原特征、原分辨率及下采样后的结果和分辨率
    else:
        return x, H, W, x, H, W  # 如果无下采样,直接返回原特征和分辨率
然后就是每个Block:SwinTransformerBlock

一组可以学习的机器(每个Swin Transformer Block就是一个这样的机器),它们会依次查看并研究这张照片上的每个小格子。

Swin Transformer Block如何一步步探索并理解图像的每个细节??

# 定义Swin Transformer Block的前向传播过程
def forward(self, x, mask_matrix):
    
    ### (1) 输入处理与规范化
    B, L, C = x.shape       # 获取批量大小B、展平的空间维度L、通道数C
    H, W = self.H, self.W   # 获取之前存储的空间分辨率H和W
    assert L == H * W, "input feature has wrong size"   # 确认展平的空间维度等于H*W
    shortcut = x    
    x = self.norm1(x)   
    x = x.view(B, H, W, C)     # 将展平的空间维度恢复为(B, H, W, C)形状

    ### (2) 对特征图进行填充,确保尺寸是窗口大小的整数倍
    # 计算需要在特征图的宽度和高度上添加的填充量,以确保尺寸是窗口大小的整数倍
    pad_l = pad_t = 0
    pad_r = (self.window_size - W % self.window_size) % self.window_size
    pad_b = (self.window_size - H % self.window_size) % self.window_size
    x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))    #  使用`F.pad`函数对特征图进行填充
    _, Hp, Wp, _ = x.shape    # 更新填充后的尺寸

    ### (3) 执行循环移位(可选)
    if self.shift_size > 0: 
        shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))  # 沿H和W维度循环移位
        attn_mask = mask_matrix  # 使用提供的mask矩阵
    else:  # 如果没有循环移位
        shifted_x = x
        attn_mask = None

    ### (4) 划分窗口
    x_windows = window_partition(shifted_x,
                                 self.window_size)  # 将移位后的特征划分成窗口,形状变为(nW*B, window_size, window_size, C)
    x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # 重新调整形状以便进行注意力计算

    ### (5) 执行窗口注意力
    attn_windows = self.attn(x_windows, mask=attn_mask)  # 在每个窗口内部执行多头自注意力机制

    #### (6) 合并窗口
    attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)  # 恢复窗口形状
    shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # 反变换回原始空间维度
    # 如果之前进行了循环移位,执行反向循环移位操作,恢复特征图的原始排列
    if self.shift_size > 0:
        x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))

    ### (7) 移除之前的填充(如果有的话)
    if pad_r > 0 or pad_b > 0:
        x = x[:, :H, :W, :].contiguous()  # 截取去除填充的部分

    ### (8) 转换回(B, H*W, C)形状
    x = x.view(B, H * W, C)

    ### (9) 进行前馈网络(FFN)操作,并应用残差连接
    x = shortcut + self.drop_path(x)  # 第一次残差连接
    x = x + self.drop_path(self.mlp(self.norm2(x)))  # 经过MLP和第二次残差连接

    # 返回前向传播的结果
    return x
  1. 窗口划分

    • 将填充并(可选)移位后的特征图划分为多个大小为(window_size, window_size)的窗口。

      # 划分窗口
      def window_partition(x, window_size):
          """
          将一个形状为 (B, H, W, C) 的四维张量划分为多个形状为 (window_size, window_size, C) 的小窗口张量,然后堆叠起来形成一个新的张量,形状为 (num_windows*B, window_size, window_size, C)
          """
          B, H, W, C = x.shape   # 获取张量维度信息,即批量大小 B、高度 H、宽度 W 和通道数 C
          x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
          windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
          return windows
      
  2. 执行窗口注意力

    • 在每个窗口内部执行多头自注意力机制(self.attn)

      WindowAttention基于窗口的多头自注意力(Window-based Multi-Head Self-Attention)模块

      它允许模型在固定的局部窗口内进行高效的自注意力计算,同时通过相对位置偏差保留了空间信息。在前向传播过程中,它首先对输入特征进行线性变换得到查询(query)、键(key)和值(value)向量,接着计算注意力分数并结合相对位置偏差,通过softmax函数归一化后,与值向量做点积得到输出特征,最后通过线性投影和丢弃层进一步处理输出

      # 前向传播过程
      def forward(self, x, mask=None):
      
          #(1) 获取输入的形状
          B_, N, C = x.shape
          
          #(2) 线性变换并分离为qkv向量
          qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
          q, k, v = qkv.unbind(0)  # 分离qkv
          # 对查询q应用缩放
          q = q * self.scale
          
          #(3) 计算注意力分数
          attn = torch.matmul(q, k.transpose(-2, -1))
      
          #(4) 添加相对位置编码
          relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
              self.window_size[0] * self.window_size[1], self.window_size[1] * self.window_size[1], -1
          )
          relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
          attn = attn + relative_position_bias.unsqueeze(0)
      
          #(5) 应用掩码并进行softmax归一化
          if mask is not None:
              nW = mask.shape[0]  # 窗口数量
              attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
              attn = attn.view(-1, self.num_heads, N, N)
              attn = self.softmax(attn)
          else:
              attn = self.softmax(attn)
          # 应用注意力权重丢弃
          attn = self.attn_drop(attn)
      
          #(6) 加权求和并进行线性投影
          x = torch.matmul(attn, v).transpose(1, 2).reshape(B_, N, C)
          x = self.proj(x)
          x = self.proj_drop(x)
          # 返回最终输出
          return x
      
  3. 窗口合并与反移位

    • 将自注意力处理后的窗口特征合并回原始空间维度,逆操作于窗口划分。

      # 窗口恢复
      def window_reverse(windows, window_size, H, W):
          """
          将一个形状为 (num_windows*B, window_size, window_size, C) 的窗口堆叠张量恢复回原始形状 (B, H, W, C),即从窗口化后的数据重构出完整数据。
          """
          # 根据窗口大小、图像高度和宽度,计算出原始数据的批量大小 B
          B = int(windows.shape[0] / (H * W / window_size / window_size))
          x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
          x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
          return x
      
    4.MLP: 实现一个具有两层隐藏层的多层感知机

前向传播过程

def forward(self, x):
    x = self.fc1(x)
    x = self.act(x)
    x = self.drop(x)
    x = self.fc2(x)
    x = self.drop(x)
    return x



## 第二部分是 sem_seg_head(`mask_former_head.py`文件)

![外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传](https://img-home.csdnimg.cn/images/20230724024159.png?origin_url=C%3A%5CUsers%5C30326%5CAppData%5CRoaming%5CTypora%5Ctypora-user-images%5Cimage-20240505114716614.png&pos_id=img-gxmMmdhw-1718015634762)

![外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传](https://img-home.csdnimg.cn/images/20230724024159.png?origin_url=C%3A%5CUsers%5C30326%5CAppData%5CRoaming%5CTypora%5Ctypora-user-images%5Cimage-20240505114825793.png&pos_id=img-QaIuy17H-1718015634763)

```python
def forward(self, features, mask=None):
 return self.layers(features, mask)

def layers(self, features, mask=None):
 
 # (1)像素解码器处理
 mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
 
 # (2)transformer 解码器处理
 # 根据transformer_in_feature 来选择如何使用上述特征调用predictor进行预测
 if self.transformer_in_feature == "multi_scale_pixel_decoder":
     predictions = self.predictor(multi_scale_features, mask_features, mask)
 else:
     if self.transformer_in_feature == "transformer_encoder":
         assert (
             transformer_encoder_features is not None
         ), "Please use the TransformerEncoderPixelDecoder."
         predictions = self.predictor(transformer_encoder_features, mask_features, mask)
     elif self.transformer_in_feature == "pixel_embedding":
         predictions = self.predictor(mask_features, mask_features, mask)
     else:
         predictions = self.predictor(features[self.transformer_in_feature], mask_features, mask)
 return predictions

核心类:pixel_decoder 和 predictor

pixel_decoder:msdeformattn.py文件

  1. 输入特征图:有四个特征图O1, O2, O3, O4,分别代表不同分辨率级别的特征表示。
  2. 处理选择:在pixel_decoder中,只对低分辨率的三层(O1, O2, O3)进行处理。这意味着O4在当前阶段未直接参与DeformAttnEncoder的处理。
  3. 像素解码器部分通过DeformAttnEncoder主要完成了低分辨率特征图的深度特征提取和跨层级的特征融合,输出编码后的特征图。
像素解码器(DeformAttnEncoder)处理流程

特征图与位置编码准备:对于O1, O2, O3,首先通过input_proj调整其维度以适应Transformer的输入要求,然后为每个特征图添加位置编码(pos),确保模型理解每个像素的相对位置信息。

	# (1)像素解码器处理
	mask_features, transformer_encoder_features, multi_scale_features = self.pixel_decoder.forward_features(features)
def forward_features(self, features):
    
    # (1)初始化srcs和pos列表:用于存储变换后的特征图和位置编码
    srcs = []  
    pos = []   
    
    # (2)特征图与位置编码处理:从低级别特征开始,反向遍历特征图(features)  
    for idx, f in enumerate(self.transformer_in_features[::-1]):
        x = features[f].float()  
        srcs.append(self.input_proj[idx](x))  # 将当前特征图x映射到适合Transformer输入的维度
        pos.append(self.pe_layer(x))          # 为特征图添加位置编码
        
    # (3)Transformer核心处理: 将处理好的特征图和位置编码传入Transformer方法
    y, spatial_shapes, level_start_index = self.transformer(srcs, pos)
  
Transformer核心处理 (forward函数)
	# (3)Transformer核心处理: srcs, pos 送入Transformer模块进行深层特征融合与提取
	y, spatial_shapes, level_start_index = self.transformer(srcs, pos)

1. 拼图碎片的整理与标记

首先,我们有三堆拼图碎片(O1;O2;O3),每堆对应于图像的一个不同分辨率层级(近景、中景和远景)。为每块碎片制作了一个透明的覆盖图,用来标记哪些部分是真正的图像信息,哪些是空白或无效的(这就是创建掩码的过程)。这样,在拼接时就能准确区分哪些是需要关注的部分。

2. 碎片的重组与层级编码

接着开始逐一处理这些碎片堆。对于每一块碎片,先把它拆成细长的拼图条,展平并调整形状。然后给每条拼图添加了一个特别的标签,上面写着“这是来自第X堆的碎片”,这样在混在一起时也能知道它原本属于哪一部分(添加层级嵌入)

3. 构建全景画卷

现在,是时候把这些分层标记好的拼图条全部汇集在一起,组成一幅大的画面。将它们按照顺序排列,从第一堆到最后一堆,确保每个碎片都紧密相连,同时保持它们之间的层级信息(使用torch.cat拼接)。这样,你就得到了一个包含所有层级信息的长条形画卷,既有宏观的整体视图,也有微观的细节展现。

4. 深入雕琢与融合

最后,你开始使用你的“魔法放大镜”——也就是编码器,来深入分析和融合这幅画卷。这个放大镜能够放大每个细节,理解每个部分与其他部分的关系,甚至能识别出不同层级间隐藏的联系(自注意力机制)。它通过一系列复杂的工艺,包括“自我反思”(self-attention),即让每个部分思考自己与整体的关系;“强化加工”(FFN,前馈网络),增强每个部分的表现力,最后,所有的信息被整合和提升,形成一幅细节丰富、层次分明的完美画卷。

def forward(self, srcs, pos_embeds):
    
    # (1)初始化与数据准备:
    # 为每个特征图创建全零mask,尺寸与特征图的高宽相同
    masks = [torch.zeros((x.size(0), x.size(2), x.size(3)), device=x.device, dtype=torch.bool) for x in srcs]
    
    # (2)逐层处理特征图:
    # 初始化用于Transformer编码器的输入变量
    src_flatten, mask_flatten, lvl_pos_embed_flatten, spatial_shapes = [], [], [], []
    
    # 遍历不同层级的特征图(srcs)、对应的位置编码(pos_embeds)以及初始化的mask
    for lvl, (src, mask, pos_embed) in enumerate(zip(srcs, masks, pos_embeds)):
        bs, c, h, w = src.shape   # 获取当前特征图的batch size、通道数、高度、宽度
        spatial_shape = (h, w)    # 记录当前特征图的空间尺寸
        spatial_shapes.append(spatial_shape)
        
        # 每个层级的特征图src、对应的mask和位置编码pos_embed被展平(将空间维度合并为一维,并转置以便序列化)
        src = src.flatten(2).transpose(1, 2) 
        mask = mask.flatten(1)
        pos_embed = pos_embed.flatten(2).transpose(1, 2)
        
        # 将位置编码与对应层级的嵌入相加,以包含层级信息
        lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)
        
        # 分别收集展平后的特征图、mask和层级位置编码
        src_flatten.append(src)
        mask_flatten.append(mask)
        lvl_pos_embed_flatten.append(lvl_pos_embed)
        
    # (3)整合输入序列
    # 将所有层级的特征图、mask和层级位置编码 使用torch.cat操作沿特定维度拼接起来,形成连续的序列,作为Transformer的输入
    src_flatten = torch.cat(src_flatten, 1)
    mask_flatten = torch.cat(mask_flatten, 1)
    lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
    
    # 将空间尺寸信息转化为Tensor,并指定设备和数据类型
    spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=src_flatten.device)
    
    # 计算每个层级在最终序列中的起始索引,用于后续处理
    level_start_index = torch.cat((spatial_shapes.new_zeros((1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))
    
    # 计算每个mask的有效比例(比如,用于处理padding区域),并堆叠为二维Tensor
    valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1)
    
    # (4) 调用编码器执行特征提取,通过自注意力等机制,融合不同层级特征,生成memory
    memory = self.encoder(src_flatten, spatial_shapes, level_start_index, valid_ratios, lvl_pos_embed_flatten, mask_flatten)
    
    # 返回编码后的特征、空间尺寸信息以及层级起始索引
    return memory, spatial_shapes, level_start_index
Encoder处理 (forward函数 - 编码器部分)
  • 编码器的核心任务是对这些序列进行处理,通过自注意力机制和潜在的前馈网络(FFN)等组件,实现特征的跨层交互和融合,输出编码后的特征memory。这个过程增强了特征表示,使之蕴含了不同尺度的空间和语义信息。
def forward(self, src, spatial_shapes, level_start_index, valid_ratios, pos=None, padding_mask=None):
    # (1)初始化输出:
    output = src  # 将输入src赋值给output,指先前编码器的输出
    # (2)获取参考点:
    reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
    # (3)遍历解码器层
    for _, layer in enumerate(self.layers):
        output = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask)

    return output

下面是是Transformer解码器中单个层(通常称为Decoder Layer)的前向传播方法:

def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None):
    # self attention
    src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)
    # 残差连接
    src = src + self.dropout1(src2)
    # 层归一化
    src = self.norm1(src)
    # ffn
    src = self.forward_ffn(src)

    return src

注意:在这一部分除了得到输出外,主要完成reference_points 的计算,帮助模型理解不同尺度空间位置的关系,这也是可变形transformer在计算时只专注于当前点周围若干点和可以跨feature map完成注意力计算的关键所在。

OK , 到这里Mask2Fomer最核心的部分就搞定了,剩下的是网络输出部分

predictor:mask2former_transformer_decoder.py文件

前向传播:用于处理多尺度特征图,并通过一个Transformer架构进行特征的编码和解码,最终输出类别预测(predictions_class)和掩码预测(predictions_mask)

比较重要的地方:

  1. 首次预测:
    • **调用forward_prediction_heads()方法,**使用初始查询特征生成初步的类别预测与掩码预测,并计算注意力掩码。
  2. Transformer解码器循环:
    • 通过循环遍历每一层Transformer解码器,交替执行交叉注意力层调用self.transformer_cross_attention_layers方法(融合源特征与查询特征)自注意力层调用self.transformer_self_attention_layers(细化查询特征间的依赖)前馈网络(FFN)
    • 每一层后,再次调用forward_prediction_heads()更新预测类别与掩码,并调整注意力掩码。
def forward(self, x, mask_features, mask = None):
    ## (1) 输入处理
    # 确保输入特征列表 `x` 的长度与模型设定的特征层级数一致
    assert len(x) == self.num_feature_levels
    
    # 初始化存储源特征(src)、位置编码(pos)及特征图尺寸(size_list)的列表
    src = []
    pos = []
    size_list = []
    
    # 由于'mask'参数未被使用,这里显式删除以避免误导或潜在错误
    del mask

    ##(2)位置编码与特征调整
    for i in range(self.num_feature_levels):   # 遍历每个特征层级
        size_list.append(x[i].shape[-2:])   # 记录当前层级特征图的尺寸到 `size_list`
        
        # 添加位置编码到 `pos`,并展平为序列形式
        pos.append(self.pe_layer(x[i], None).flatten(2))
        
        # 对特征图应用输入投影层,并加入层级嵌入,结果存入 `src`
        src.append(self.input_proj[i](x[i]).flatten(2) + self.level_embed.weight[i][None, :, None])
        
        # 调整 `pos` 和 `src` 的维度顺序为(HW, N, C),适配Transformer输入格式
        pos[-1] = pos[-1].permute(2, 0, 1)
        src[-1] = src[-1].permute(2, 0, 1)

    # 获取第一个特征层级的批量大小和通道数信息
    _, bs, _ = src[0].shape

    ## (3) 初始化查询嵌入与输出特征
    # 重复查询嵌入和特征以匹配批次大小,准备Transformer解码器输入
    query_embed = self.query_embed.weight.unsqueeze(1).repeat(1, bs, 1)
    output = self.query_feat.weight.unsqueeze(1).repeat(1, bs, 1)

    
    ##(4)首次预测
    # 初始化预测类别和掩码的列表
    predictions_class = []
    predictions_mask = []
    # 使调用forward_prediction_heads()方法,使用初始查询特征生成初步的类别预测与掩码预测,并计算注意力掩码。
    outputs_class, outputs_mask, attn_mask = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0])
    predictions_class.append(outputs_class)
    predictions_mask.append(outputs_mask)

    ##(5)Transformer层循环*重点
    for i in range(self.num_layers):  # 遍历Transformer的每一层
        level_index = i % self.num_feature_levels  # 循环选择特征层级
        
        # 更新注意力掩码,避免全为True的情况
        attn_mask[torch.where(attn_mask.sum(-1) == attn_mask.shape[-1])] = False
        
        # 交叉注意力层:融合源特征和查询特征 
        output = self.transformer_cross_attention_layers[i](
            output, src[level_index],
            memory_mask=attn_mask,
            memory_key_padding_mask=None,  # 不对padding区域应用mask
            pos=pos[level_index], query_pos=query_embed
        )
        
        # 自注意力层:在查询特征间细化特征表示
        output = self.transformer_self_attention_layers[i](
            output, tgt_mask=None,
            tgt_key_padding_mask=None,
            query_pos=query_embed
        )
        
        # 前馈网络(FFN):进一步加工输出
        output = self.transformer_ffn_layers[i](output)
        
        # 每层后进行预测并更新预测列表
        outputs_class, outputs_mask, attn_mask = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[(i + 1) % self.num_feature_levels])
        predictions_class.append(outputs_class)
        predictions_mask.append(outputs_mask)

    # 校验预测类别列表长度,确保与Transformer层数相符
    assert len(predictions_class) == self.num_layers + 1

	## (6) 构造输出字典
    out = {
        'pred_logits': predictions_class[-1],  # 最终的类别预测
        'pred_masks': predictions_mask[-1],    # 最终的掩码预测
        'aux_outputs': self._set_aux_loss(     # 辅助输出处理
            predictions_class if self.mask_classification else None, predictions_mask 
        )
    }
    return out
(1)forward_prediction_heads()方法: 使用学习到的查询特征(output)进行首次预测
## 使用forward_prediction_heads()方法 利用学习到的查询特征(output)进行首次预测
    outputs_class, outputs_mask, attn_mask = self.forward_prediction_heads(output, mask_features, attn_mask_target_size=size_list[0])
    predictions_class.append(outputs_class)
    predictions_mask.append(outputs_mask)
def forward_prediction_heads(self, output, mask_features, attn_mask_target_size):
    
    ##(1)标准化(Normalization)与转置:
    decoder_output = self.decoder_norm(output)    # 对 output 进行LayerNorm标准化
    decoder_output = decoder_output.transpose(0, 1)  # 将解码器输出的形状转换为 (sequence_length, batch_size, feature_dim)

    ##(2)分类与掩码特征生成:
    # 应用class_embed层将标准化后的特征映射到类别空间,生成每个位置的类别预测outputs_class
    outputs_class = self.class_embed(decoder_output)  
    
     # 通过mask_embed层处理相同特征,得到用于生成掩码的特征表示mask_embed
    mask_embed = self.mask_embed(decoder_output)
     # 利用掩码特征(mask_embed)与原始特征图(mask_features)的矩阵运算,初步生成掩码预测
    outputs_mask = torch.einsum("bqc,bchw->bqhw", mask_embed, mask_features)

	##(3)注意力掩码(Attention Mask)处理:
    # 使用双线性插值调整掩码尺寸至目标大小(attn_mask_target_size)
    attn_mask = F.interpolate(outputs_mask, size=attn_mask_target_size, mode="bilinear", align_corners=False)
    
    # 调整尺寸后的掩码通过sigmoid函数转换为概率值,并设定阈值(0.5)以二值化处理,得到布尔型注意力掩码
    attn_mask = (attn_mask.sigmoid().flatten(2).unsqueeze(1).repeat(1, self.num_heads, 1, 1).flatten(0, 1) < 0.5).bool()
    
    # 将生成的注意力掩码设置为不可训练(detach),意味着其梯度不会在反向传播中计算
    attn_mask = attn_mask.detach()

    # 返回分类预测、掩码预测以及处理过的注意力掩码
    return outputs_class, outputs_mask, attn_mask

(2)在Transformer解码器的每层中transformer_cross_attention_layers方法通过以下步骤融合源(编码器)特征与查询(解码器)特征融合源特征和查询特征
  • 核心操作
    • 位置编码应用:首先,解码器的输出output和编码器的特征src[level_index]各自加上对应位置的嵌入query_pospos[level_index],以融入序列中元素的位置信息。
    • 注意力计算:接下来,利用多头注意力机制计算解码器输出(作为查询)与编码器特征(作为键和值)之间的关系。
    • 信息融合与更新:依据计算出的注意力权重,更新解码器输出。这包括残差连接(原输出与注意力结果相加)及潜在的Dropout应用
	# 在解码器的每一层中,此行代码执行交叉注意力操作
	output = self.transformer_cross_attention_layers[i](
    	output, src[level_index],      # 输出(解码器当前层的输出)和来源特征(编码器对应层的输出)
        memory_mask=attn_mask,         # 注意力掩码,用于限制哪些部分可被注意力机制访问
        memory_key_padding_mask=None,  
    	pos=pos[level_index],          # 来源序列的位置嵌入
        query_pos=query_embed          # 查询(解码器)的位置嵌入
	)
## Transformer模型解码器中交叉注意力层的前向传播过程:通过交叉注意力机制整合解码器的查询信息与编码器的特征
def forward(self, tgt, memory,          # tgt是解码器输入,memory是编码器输出
            memory_mask=None,           # 编码器输出的注意力掩码
            memory_key_padding_mask=None, # 编码器输出的padding掩码,默认不使用
            pos=None,                   # 位置编码向量,用于源序列
            query_pos=None):             # 位置编码向量,用于查询序列(解码器侧)
    # 先进行LN,然后执行注意力操作,适用于需要稳定训练过程的情形。
    if self.normalize_before:
        return self.forward_pre(tgt, memory, memory_mask,
                                memory_key_padding_mask, pos, query_pos)
    # 先执行注意力操作,之后进行LN,可能对模型收敛速度有益
    return self.forward_post(tgt, memory, memory_mask,
                             memory_key_padding_mask, pos, query_pos)

其中self.forward_pre()方法:先进行LN,然后执行注意力操作

def forward_pre(self, tgt, memory, 
                memory_mask=None, 
                memory_key_padding_mask=None, 
                pos=None, 
                query_pos=None):
    # (1)对解码器的输入(tgt)应用LayerNorm
    tgt2 = self.norm(tgt)
    
    # (2)执行多头注意力操作,结合位置嵌入
    tgt2 = self.multihead_attn(
        query=self.with_pos_embed(tgt2, query_pos),  # 添加查询位置嵌入
        key=self.with_pos_embed(memory, pos),        # 添加记忆位置嵌入
        value=memory,                                # 使用编码器的输出作为价值信息
        attn_mask=memory_mask,                       # 应用注意力掩码
        key_padding_mask=memory_key_padding_mask     # 默认不使用padding掩码
    )[0]  # multihead_attn返回一个元组,我们只取第一项(注意力后的输出)
    
    # (3)残差连接后应用dropout
    tgt = tgt + self.dropout(tgt2)
    
    # 返回更新后的解码器输出
    return tgt  

self.forward_post()方法:先执行注意力操作,之后进行LN

def forward_post(self, tgt, memory, 
                 memory_mask=None, 
                 memory_key_padding_mask=None, 
                 pos=None, 
                 query_pos=None):
    
    # (1)直接执行多头注意力操作,不先做LayerNorm
    tgt2 = self.multihead_attn(
        query=self.with_pos_embed(tgt, query_pos),
        key=self.with_pos_embed(memory, pos),
        value=memory,
        attn_mask=memory_mask,
        key_padding_mask=memory_key_padding_mask
    )[0]
    
    # (2)残差连接后应用dropout
    tgt = tgt + self.dropout(tgt2)
    
    # (3)然后应用LayerNorm
    tgt = self.norm(tgt)
    
    # 返回最终更新后的解码器输出
    return tgt  
(3)自注意力:对output(查询特征)应用自注意力机制,目的是捕捉特征内部的依赖关系,从而精细化特征表示
##  自注意力:在查询特征间执行自注意力操作,进一步细化查询特征表示
    output = self.transformer_self_attention_layers[i](
        output,                    # 查询特征
        tgt_mask=None,
        tgt_key_padding_mask=None,
        query_pos=query_embed      # 位置编码
    )
def forward(self, tgt,
            tgt_mask: Optional[Tensor] = None,
            tgt_key_padding_mask: Optional[Tensor] = None,
            query_pos: Optional[Tensor] = None):
    
    # 先进行LN,然后执行注意力操作
    if self.normalize_before:
        return self.forward_pre(tgt, tgt_mask,
                                tgt_key_padding_mask, query_pos)
    # 先执行注意力操作,之后进行LN
    return self.forward_post(tgt, tgt_mask,
                             tgt_key_padding_mask, query_pos)

其中self.forward_pre(): 先进行LN,然后执行注意力操作

def forward_pre(self, tgt,
                tgt_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):

    # (1)对输入tgt进行Layer Normalization
    tgt2 = self.norm(tgt)
    
    # (2)将Layer Norm后的tgt与位置编码(query_pos)相结合,准备作为自注意力的查询和键
    q = k = self.with_pos_embed(tgt2, query_pos)
    
    # (3)执行自注意力计算,tgt2作为自注意力的值
    tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
    
    # (4)自注意力的结果与原始tgt残差连接,并应用dropout
    tgt = tgt + self.dropout(tgt2)
    
    # 返回更新后的tgt
    return tgt

self.forward_post(): 先执行注意力操作,之后进行LN

def forward_post(self, tgt,
                 tgt_mask: Optional[Tensor] = None,
                 tgt_key_padding_mask: Optional[Tensor] = None,
                 query_pos: Optional[Tensor] = None):
    
    # (1)直接将tgt与位置编码(query_pos)相结合,准备作为自注意力的查询和键
    q = k = self.with_pos_embed(tgt, query_pos)
    
    # (2)执行自注意力计算,注意此时直接使用tgt作为自注意力的值
    tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
    
    # (3)自注意力的结果与原始tgt残差连接,并应用dropout
    tgt = tgt + self.dropout(tgt2)
    
    # (4)在所有操作之后执行Layer Normalization
    tgt = self.norm(tgt)
    
    # 返回更新后的tgt
    return tgt

的值
tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]

# (4)自注意力的结果与原始tgt残差连接,并应用dropout
tgt = tgt + self.dropout(tgt2)

# 返回更新后的tgt
return tgt

**self.forward_post(): 先执行注意力操作,之后进行LN**

```python
def forward_post(self, tgt,
                 tgt_mask: Optional[Tensor] = None,
                 tgt_key_padding_mask: Optional[Tensor] = None,
                 query_pos: Optional[Tensor] = None):
    
    # (1)直接将tgt与位置编码(query_pos)相结合,准备作为自注意力的查询和键
    q = k = self.with_pos_embed(tgt, query_pos)
    
    # (2)执行自注意力计算,注意此时直接使用tgt作为自注意力的值
    tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
    
    # (3)自注意力的结果与原始tgt残差连接,并应用dropout
    tgt = tgt + self.dropout(tgt2)
    
    # (4)在所有操作之后执行Layer Normalization
    tgt = self.norm(tgt)
    
    # 返回更新后的tgt
    return tgt
  • 24
    点赞
  • 33
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值