Restormer自定义训练测试代码，注释详尽适合学习

## Restormer: Efficient Transformer for High-Resolution Image Restoration ## Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang ## https://arxiv.org/abs/2111.09881 import torch import torch.nn as nn import torch.nn.functional as F from pdb import set_trace as stx import numbers from einops import rearrange ########################################################################## ## Layer Norm def to_3d(x): return rearrange(x, 'b c h w -> b (h w) c') def to_4d(x, h, w): return rearrange(x, 'b (h w) c -> b c h w', h=h, w=w) class BiasFree_LayerNorm(nn.Module): def __init__(self, normalized_shape): super(BiasFree_LayerNorm, self).__init__() if isinstance(normalized_shape, numbers.Integral): normalized_shape = (normalized_shape,) normalized_shape = torch.Size(normalized_shape) assert len(normalized_shape) == 1 self.weight = nn.Parameter(torch.ones(normalized_shape)) self.normalized_shape = normalized_shape def forward(self, x): sigma = x.var(-1, keepdim=True, unbiased=False) return x / torch.sqrt(sigma + 1e-5) * self.weight class WithBias_LayerNorm(nn.Module): def __init__(self, normalized_shape): super(WithBias_LayerNorm, self).__init__() if isinstance(normalized_shape, numbers.Integral): normalized_shape = (normalized_shape,) normalized_shape = torch.Size(normalized_shape) assert len(normalized_shape) == 1 self.weight = nn.Parameter(torch.ones(normalized_shape)) self.bias = nn.Parameter(torch.zeros(normalized_shape)) self.normalized_shape = normalized_shape def forward(self, x): mu = x.mean(-1, keepdim=True) sigma = x.var(-1, keepdim=True, unbiased=False) return (x - mu) / torch.sqrt(sigma + 1e-5) * self.weight + self.bias class LayerNorm(nn.Module): def __init__(self, dim, LayerNorm_type): super(LayerNorm, self).__init__() if LayerNorm_type == 'BiasFree': self.body = BiasFree_LayerNorm(dim) else: self.body = WithBias_LayerNorm(dim) def forward(self, x): h, w = x.shape[-2:] return to_4d(self.body(to_3d(x)), h, w) ########################################################################## ## Gated-Dconv Feed-Forward Network (GDFN) class FeedForward(nn.Module): def __init__(self, dim, ffn_expansion_factor, bias): super(FeedForward, self).__init__() # 隐藏层特征维度等于输入维度乘以扩张因子 hidden_features = int(dim * ffn_expansion_factor) # 1*1 升维 self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias) # 3*3 分组卷积 self.dwconv = nn.Conv2d(hidden_features * 2, hidden_features * 2, kernel_size=3, stride=1, padding=1, groups=hidden_features * 2, bias=bias) # 1*1 降维 self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias) def forward(self, x): x = self.project_in(x) x1, x2 = self.dwconv(x).chunk(2, dim=1) # 第 1 个维度方向切分成 2 块 x = F.gelu(x1) * x2 # gelu 相当于 relu+dropout x = self.project_out(x) return x ########################################################################## ## Multi-DConv Head Transposed Self-Attention (MDTA) class Attention(nn.Module): def __init__(self, dim, num_heads, bias): super(Attention, self).__init__() self.num_heads = num_heads # 注意力头的个数 self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1)) # 可学习系数 # 1*1 升维 self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias) # 3*3 分组卷积 self.qkv_dwconv = nn.Conv2d(dim * 3, dim * 3, kernel_size=3, stride=1, padding=1, groups=dim * 3, bias=bias) # 1*1 卷积 self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias) def forward(self, x): b, c, h, w = x.shape # 输入的结构 batch 数，通道数和高宽 qkv = self.qkv_dwconv(self.qkv(x)) q, k, v = qkv.chunk(3, dim=1) # 第 1 个维度方向切分成 3 块 # 改变 q, k, v 的结构为 b head c (h w)，将每个二维 plane 展平 q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads) k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads) v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads) q = torch.nn.functional.normalize(q, dim=-1) # C 维度标准化，这里的 C 与通道维度略有不同 k = torch.nn.functional.normalize(k, dim=-1) attn = (q @ k.transpose(-2, -1)) * self.temperature attn = attn.softmax(dim=-1) out = (attn @ v) # 注意力图(严格来说不算图) # 将展平后的注意力图恢复 out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w) # 真正的注意力图 out = self.project_out(out) return out ########################################################################## class TransformerBlock(nn.Module): def __init__(self, dim, num_heads, ffn_expansion_factor, bias, LayerNorm_type): super(TransformerBlock, self).__init__() self.norm1 = LayerNorm(dim, LayerNorm_type) # 层标准化 self.attn = Attention(dim, num_heads, bias) # 自注意力 self.norm2 = LayerNorm(dim, LayerNorm_type) # 层表转化 self.ffn = FeedForward(dim, ffn_expansion_factor, bias) # FFN def forward(self, x): x = x + self.attn(self.norm1(x)) # 残差 x = x + self.ffn(self.norm2(x)) # 残差 return x ########################################################################## ## Overlapped image patch embedding with 3x3 Conv class OverlapPatchEmbed(nn.Module): def __init__(self, in_c=3, embed_dim=48, bias=False): super(OverlapPatchEmbed, self).__init__() self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=3, stride=1, padding=1, bias=bias) def forward(self, x): x = self.proj(x) return x ########################################################################## ## Resizing modules class Downsample(nn.Module): # 下采样 PixelUnshuffle def __init__(self, n_feat): super(Downsample, self).__init__() self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat // 2, kernel_size=3, stride=1, padding=1, bias=False), nn.PixelUnshuffle(2)) def forward(self, x): return self.body(x) class Upsample(nn.Module): # 上采样 PixelShuffle def __init__(self, n_feat): super(Upsample, self).__init__() self.body = nn.Sequential(nn.Conv2d(n_feat, n_feat * 2, kernel_size=3, stride=1, padding=1, bias=False), nn.PixelShuffle(2)) def forward(self, x): return self.body(x) ########################################################################## ##---------- Restormer ----------------------- class Restormer(nn.Module): def __init__(self, inp_channels=3, out_channels=3, dim=48, # 特征图维度 num_blocks=[4, 6, 6, 8], num_refinement_blocks=4, heads=[1, 2, 4, 8], ffn_expansion_factor=2.66, # 扩展因子 bias=False, LayerNorm_type='WithBias', ## Other option 'BiasFree' dual_pixel_task=False ## True for dual-pixel defocus deblurring only. Also set inp_channels=6 ):