【论文速看】DL最新进展202411011-图像超分、Transformer

### 图像超分辨率 Transformer 模型实现代码 以下是基于 SRFormer 的图像超分辨率 Transformer 架构的一个简化版本实现。该模型通过引入 PSA（Position-Specific Attention）机制来提升性能，同时保持了良好的可扩展性和易用性。 #### 1. 定义核心模块：PSA 和 Transformer 层 ```python import torch import torch.nn as nn class PositionSpecificAttention(nn.Module): """位置特定注意力 (PSA) 模块""" def __init__(self, channels, reduction=4): super(PositionSpecificAttention, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(channels, channels // reduction), nn.ReLU(inplace=True), nn.Linear(channels // reduction, channels), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) y = self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x) class TransformerLayer(nn.Module): """Transformer 层定义""" def __init__(self, dim, num_heads=8, dropout=0.1): super(TransformerLayer, self).__init__() self.self_attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout) self.norm1 = nn.LayerNorm(dim) self.ffn = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim * 4, dim) ) self.norm2 = nn.LayerNorm(dim) def forward(self, x): attn_output, _ = self.self_attn(x, x, x) x = x + self.norm1(attn_output) ffn_output = self.ffn(x) x = x + self.norm2(ffn_output) return x ``` #### 2. 整合 SRFormer 主体结构 ```python class SRFormer(nn.Module): """SRFormer 模型主体""" def __init__(self, in_channels=3, out_channels=3, embed_dim=64, depth=6): super(SRFormer, self).__init__() self.conv_first = nn.Conv2d(in_channels, embed_dim, kernel_size=3, stride=1, padding=1) transformer_layers = [] for i in range(depth): transformer_layers.append(TransformerLayer(embed_dim)) self.transformer = nn.Sequential(*transformer_layers) self.psa = PositionSpecificAttention(channels=embed_dim) self.reconstruction = nn.Sequential( nn.Conv2d(embed_dim, embed_dim, kernel_size=3, stride=1, padding=1), nn.LeakyReLU(negative_slope=0.2, inplace=True), nn.Conv2d(embed_dim, out_channels, kernel_size=3, stride=1, padding=1) ) def forward(self, x): feat = self.conv_first(x) h, w = feat.shape[-2:] feat = feat.flatten(2).permute(2, 0, 1) # BxCxHxW -> LxBxC feat = self.transformer(feat) feat = feat.permute(1, 2, 0).view(-1, feat.size(2), h, w) # LxBxC -> BxCxHxW feat = self.psa(feat) output = self.reconstruction(feat) return output ``` #### 3. 使用示例 ```python if __name__ == "__main__": model = SRFormer(in_channels=3, out_channels=3, embed_dim=64, depth=6) input_tensor = torch.randn(1, 3, 64, 64) # 输入低分辨率图像 output_tensor = model(input_tensor) # 输出高分辨率图像 print("Input Shape:", input_tensor.shape) print("Output Shape:", output_tensor.shape) ``` --- ### 关键技术解析 上述代码实现了 SRFormer 中的核心组件，包括： - **PSA 模块**：通过自适应权重调整不同通道的重要性，从而增强局部特征表达能力[^2]。 - **Transformer Layer**：利用多头注意力机制捕获全局依赖关系，并结合前馈网络完成特征变换。 - **Residual Connection & Normalization**：在每一层加入残差连接和归一化操作，有助于缓解梯度消失问题并加速收敛。 此外，在实际应用中可以进一步优化采样策略或增加额外损失函数以提高重建质量[^1]。 ---