【Block总结】掩码窗口自注意力 (M-WSA)

在这里插入图片描述

摘要

论文链接：https://arxiv.org/pdf/2404.07846
论文标题：Transformer-Based Blind-Spot Network for Self-Supervised Image Denoising
Masked Window-Based Self-Attention (M-WSA) 是一种新颖的自注意力机制，旨在解决传统自注意力方法在处理图像时的局限性，特别是在图像去噪和恢复任务中。M-WSA 通过引入掩码机制，确保在计算注意力时遵循盲点要求，从而避免信息泄露。

设计原理

窗口自注意力：M-WSA 基于窗口自注意力（Window Self-Attention, WSA）的概念，将输入图像划分为多个不重叠的窗口。在每个窗口内，计算自注意力以捕捉局部特征。这种方法的计算复杂度相对较低，适合处理高分辨率图像。
掩码机制：为了满足盲点要求，M-WSA 在计算注意力时应用了掩码。具体而言，掩码限制了每个像素只能关注其窗口内的特定像素，从而避免了对盲点信息的访问。这一设计确保了网络在去噪时不会泄露噪声信息。
扩张卷积模拟：M-WSA 的掩码设计模仿了扩张卷积的感受野，使得网络能够在保持计算效率的同时，捕捉到更大范围的上下文信息。这种方法有效地扩展了网络的感受野，增强了特征提取能力。

优势

高效性：通过限制注意力计算在窗口内，M-WSA 显著降低了计算复杂度，使其适用于大规模图像处理任务。
信息保护：掩码机制确保了盲点信息不被泄露，从而提高了去噪效果，特别是在处理具有空间相关噪声的图像时。
灵活性：M-WSA 可以与其他网络架构结合使用，增强其在各种视觉任务中的表现，尤其是在自我监督学习和图像恢复领域。

实验结果

在多个真实世界的图像去噪数据集上进行的实验表明，M-WSA 显著提高了去噪性能，超越了传统的卷积网络和其他自注意力机制。这一结果表明，M-WSA 在处理复杂噪声模式时具有良好的适应性和有效性。

代码

Masked Window-Based Self-Attention (M-WSA) 通过结合窗口自注意力和掩码机制，为图像去噪和恢复任务提供了一种有效的解决方案。其设计不仅提高了计算效率，还确保了信息的安全性，展示了在自我监督学习中的广泛应用潜力。代码：

import torch
import torch.nn as nn
from einops import rearrange
from torch import einsumdef to(x):return {'device': x.device, 'dtype': x.dtype}def expand_dim(t, dim, k):t = t.unsqueeze(dim=dim)expand_shape = [-1] * len(t.shape)expand_shape[dim] = kreturn t.expand(*expand_shape)def rel_to_abs(x):b, l, m = x.shaper = (m + 1) // 2col_pad = torch.zeros((b, l, 1), **to(x))x = torch.cat((x, col_pad), dim=2)flat_x = rearrange(x, 'b l c -> b (l c)')flat_pad = torch.zeros((b, m - l), **to(x))flat_x_padded = torch.cat((flat_x, flat_pad), dim=1)final_x = flat_x_padded.reshape(b, l + 1, m)final_x = final_x[:, :l, -r:]return final_xdef relative_logits_1d(q, rel_k):b, h, w, _ = q.shaper = (rel_k.shape[0] + 1) // 2logits = einsum('b x y d, r d -> b x y r', q, rel_k)logits = rearrange(logits, 'b x y r -> (b x) y r')logits = rel_to_abs(logits)logits = logits.reshape(b, h, w, r)logits = expand_dim(logits, dim=2, k=r)return logitsclass RelPosEmb(nn.Module):def __init__(self,block_size,rel_size,dim_head):super().__init__()height = width = rel_sizescale = dim_head ** -0.5self.block_size = block_sizeself.rel_height = nn.Parameter(torch.randn(height * 2 - 1, dim_head) * scale)self.rel_width = nn.Parameter(torch.randn(width * 2 - 1, dim_head) * scale)def forward(self, q):block = self.block_sizeq = rearrange(q, 'b (x y) c -> b x y c', x=block)rel_logits_w = relative_logits_1d(q, self.rel_width)rel_logits_w = rearrange(rel_logits_w, 'b x i y j-> b (x y) (i j)')q = rearrange(q, 'b x y d -> b y x d')rel_logits_h = relative_logits_1d(q, self.rel_height)rel_logits_h = rearrange(rel_logits_h, 'b x i y j -> b (y x) (j i)')return rel_logits_w + rel_logits_hclass FixedPosEmb(nn.Module):def __init__(self, window_size, overlap_window_size):super().__init__()self.window_size = window_sizeself.overlap_window_size = overlap_window_sizeattention_mask_table = torch.zeros((window_size + overlap_window_size - 1),(window_size + overlap_window_size - 1))attention_mask_table[0::2, :] = float('-inf')attention_mask_table[:, 0::2] = float('-inf')attention_mask_table = attention_mask_table.view((window_size + overlap_window_size - 1) * (window_size + overlap_window_size - 1))# get pair-wise relative position index for each token inside the windowcoords_h = torch.arange(self.window_size)coords_w = torch.arange(self.window_size)coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Wwcoords_flatten_1 = torch.flatten(coords, 1)  # 2, Wh*Wwcoords_h = torch.arange(self.overlap_window_size)coords_w = torch.arange(self.overlap_window_size)coords = torch.stack(torch.meshgrid([coords_h, coords_w]))coords_flatten_2 = torch.flatten(coords, 1)relative_coords = coords_flatten_1[:, :, None] - coords_flatten_2[:, None, :]  # 2, Wh*Ww, Wh*Wwrelative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2relative_coords[:, :, 0] += self.overlap_window_size - 1  # shift to start from 0relative_coords[:, :, 1] += self.overlap_window_size - 1relative_coords[:, :, 0] *= self.window_size + self.overlap_window_size - 1relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Wwself.attention_mask = nn.Parameter(attention_mask_table[relative_position_index.view(-1)].view(1, self.window_size ** 2, self.overlap_window_size ** 2), requires_grad=False)def forward(self):return self.attention_maskclass DilatedOCA(nn.Module):def __init__(self, dim, window_size, overlap_ratio, num_heads, dim_head, bias):super(DilatedOCA, self).__init__()self.num_spatial_heads = num_headsself.dim = dimself.window_size = window_sizeself.overlap_win_size = int(window_size * overlap_ratio) + window_sizeself.dim_head = dim_headself.inner_dim = self.dim_head * self.num_spatial_headsself.scale = self.dim_head ** -0.5self.unfold = nn.Unfold(kernel_size=(self.overlap_win_size, self.overlap_win_size), stride=window_size,padding=(self.overlap_win_size - window_size) // 2)self.qkv = nn.Conv2d(self.dim, self.inner_dim * 3, kernel_size=1, bias=bias)self.project_out = nn.Conv2d(self.inner_dim, dim, kernel_size=1, bias=bias)self.rel_pos_emb = RelPosEmb(block_size=window_size,rel_size=window_size + (self.overlap_win_size - window_size),dim_head=self.dim_head)self.fixed_pos_emb = FixedPosEmb(window_size, self.overlap_win_size)def forward(self, x):b, c, h, w = x.shapeqkv = self.qkv(x)qs, ks, vs = qkv.chunk(3, dim=1)# spatial attentionqs = rearrange(qs, 'b c (h p1) (w p2) -> (b h w) (p1 p2) c', p1=self.window_size, p2=self.window_size)ks, vs = map(lambda t: self.unfold(t), (ks, vs))ks, vs = map(lambda t: rearrange(t, 'b (c j) i -> (b i) j c', c=self.inner_dim), (ks, vs))# print(f'qs.shape:{qs.shape}, ks.shape:{ks.shape}, vs.shape:{vs.shape}')# split headsqs, ks, vs = map(lambda t: rearrange(t, 'b n (head c) -> (b head) n c', head=self.num_spatial_heads),(qs, ks, vs))# attentionqs = qs * self.scalespatial_attn = (qs @ ks.transpose(-2, -1))spatial_attn += self.rel_pos_emb(qs)spatial_attn += self.fixed_pos_emb()spatial_attn = spatial_attn.softmax(dim=-1)out = (spatial_attn @ vs)out = rearrange(out, '(b h w head) (p1 p2) c -> b (head c) (h p1) (w p2)', head=self.num_spatial_heads,h=h // self.window_size, w=w // self.window_size, p1=self.window_size, p2=self.window_size)# merge spatial and channelout = self.project_out(out)return outif __name__ == "__main__":dim = 64window_size = 8overlap_ratio = 0.5num_heads = 2dim_head = 16# 初始化 DilatedOCA 模块oca_attention = DilatedOCA(dim=dim,window_size=window_size,overlap_ratio=overlap_ratio,num_heads=num_heads,dim_head=dim_head,bias=True)device = torch.device("cuda" if torch.cuda.is_available() else "cpu")oca_attention = oca_attention.to(device)print(oca_attention)x = torch.randn(1, 32, 640, 480).to(device)# 前向传播output = oca_attention(x)print("input张量形状:", x.shape)print("output张量形状:", output.shape)

DilatedOCA模块详解

代码结构

import torch
import torch.nn as nn
from einops import rearrange

导入库：首先导入 PyTorch 和 einops 库。einops 用于简化张量的重排操作。

模块定义

class DilatedOCA(nn.Module):def __init__(self, dim, window_size, overlap_ratio, num_heads, dim_head, bias):super(DilatedOCA, self).__init__()self.num_spatial_heads = num_headsself.dim = dimself.window_size = window_sizeself.overlap_win_size = int(window_size * overlap_ratio) + window_sizeself.dim_head = dim_headself.inner_dim = self.dim_head * self.num_spatial_headsself.scale = self.dim_head ** -0.5self.unfold = nn.Unfold(kernel_size=(self.overlap_win_size, self.overlap_win_size), stride=window_size,padding=(self.overlap_win_size - window_size) // 2)self.qkv = nn.Conv2d(self.dim, self.inner_dim * 3, kernel_size=1, bias=bias)self.project_out = nn.Conv2d(self.inner_dim, dim, kernel_size=1, bias=bias)self.rel_pos_emb = RelPosEmb(block_size=window_size,rel_size=window_size + (self.overlap_win_size - window_size),dim_head=self.dim_head)self.fixed_pos_emb = FixedPosEmb(window_size, self.overlap_win_size)

初始化方法：__init__ 方法定义了模块的结构。
- dim：输入特征的通道数。
- window_size：窗口的大小，用于空间注意力计算。
- overlap_ratio：重叠窗口的比例，决定了窗口之间的重叠程度。
- num_heads：空间注意力的头数。
- dim_head：每个头的维度。
层的定义：
- self.unfold：用于将输入张量展开为重叠窗口的操作。
- self.qkv：一个 1x1 的卷积层，用于生成查询（Q）、键（K）和值（V）三个特征图。
- self.project_out：一个 1x1 的卷积层，用于将输出特征映射回原始通道数。
- self.rel_pos_emb 和 self.fixed_pos_emb：用于位置编码的模块，增强模型对空间位置的感知。

前向传播

def forward(self, x):b, c, h, w = x.shapeqkv = self.qkv(x)qs, ks, vs = qkv.chunk(3, dim=1)# spatial attentionqs = rearrange(qs, 'b c (h p1) (w p2) -> (b h w) (p1 p2) c', p1=self.window_size, p2=self.window_size)ks, vs = map(lambda t: self.unfold(t), (ks, vs))ks, vs = map(lambda t: rearrange(t, 'b (c j) i -> (b i) j c', c=self.inner_dim), (ks, vs))# split headsqs, ks, vs = map(lambda t: rearrange(t, 'b n (head c) -> (b head) n c', head=self.num_spatial_heads),(qs, ks, vs))# attentionqs = qs * self.scalespatial_attn = (qs @ ks.transpose(-2, -1))spatial_attn += self.rel_pos_emb(qs)spatial_attn += self.fixed_pos_emb()spatial_attn = spatial_attn.softmax(dim=-1)out = (spatial_attn @ vs)out = rearrange(out, '(b h w head) (p1 p2) c -> b (head c) (h p1) (w p2)', head=self.num_spatial_heads,h=h // self.window_size, w=w // self.window_size, p1=self.window_size, p2=self.window_size)# merge spatial and channelout = self.project_out(out)return out

输入形状：x 的形状为 (batch_size, channels, height, width)，其中 b 是批量大小，c 是通道数，h 和 w 是图像的高度和宽度。
特征提取：
- qkv = self.qkv(x)：通过 qkv 层生成 Q、K、V 特征图。
- qs, ks, vs = qkv.chunk(3, dim=1)：将 Q、K、V 特征图沿通道维度分离。
空间注意力计算：
- qs 被重排为适合空间注意力计算的格式。
- ks 和 vs 通过 unfold 操作展开为重叠窗口。
分头处理：
- 使用 einops.rearrange 将 Q、K、V 的形状调整为适合多头自注意力计算的格式。
计算注意力：
- qs = qs * self.scale：对 Q 进行缩放以提高稳定性。
- spatial_attn = (qs @ ks.transpose(-2, -1))：计算注意力分数。
- spatial_attn += self.rel_pos_emb(qs) 和 spatial_attn += self.fixed_pos_emb()：添加位置编码以增强空间感知。
- spatial_attn = spatial_attn.softmax(dim=-1)：对注意力分数进行 softmax 归一化。
输出计算：
- out = (spatial_attn @ vs)：使用注意力权重对 V 进行加权求和，得到最终输出。
重排输出：
- out = rearrange(out, '(b h w head) (p1 p2) c -> b (head c) (h p1) (w p2)', ...)：将输出重排回原始形状。
最终投影：
- out = self.project_out(out)：通过投影层将输出映射回原始通道数。