""" SwiftFormerTemporal: Temporal extension of SwiftFormer for frame prediction """ import torch import torch.nn as nn from .swiftformer import ( SwiftFormer, SwiftFormer_depth, SwiftFormer_width, stem, Embedding, Stage ) from timm.layers import DropPath, trunc_normal_ class DecoderBlock(nn.Module): """Upsampling block for frame prediction decoder with residual connections""" def __init__(self, in_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1): super().__init__() # 主路径：反卷积 + 两个卷积层 self.conv_transpose = nn.ConvTranspose2d( in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, output_padding=output_padding, bias=True # 启用bias，因为移除了BN ) self.conv1 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=True) self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=True) # 残差路径：如果需要改变通道数或空间尺寸 self.shortcut = nn.Identity() if in_channels != out_channels or stride != 1: # 使用1x1卷积调整通道数，如果需要上采样则使用反卷积 if stride == 1: self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True) else: self.shortcut = nn.ConvTranspose2d( in_channels, out_channels, kernel_size=1, stride=stride, padding=0, output_padding=output_padding, bias=True ) # 使用LeakyReLU避免死亡神经元 self.activation = nn.LeakyReLU(0.2, inplace=True) # 初始化权重 self._init_weights() def _init_weights(self): # 初始化反卷积层 nn.init.kaiming_normal_(self.conv_transpose.weight, mode='fan_out', nonlinearity='leaky_relu') if self.conv_transpose.bias is not None: nn.init.constant_(self.conv_transpose.bias, 0) # 初始化卷积层 nn.init.kaiming_normal_(self.conv1.weight, mode='fan_out', nonlinearity='leaky_relu') if self.conv1.bias is not None: nn.init.constant_(self.conv1.bias, 0) nn.init.kaiming_normal_(self.conv2.weight, mode='fan_out', nonlinearity='leaky_relu') if self.conv2.bias is not None: nn.init.constant_(self.conv2.bias, 0) # 初始化shortcut if not isinstance(self.shortcut, nn.Identity): if isinstance(self.shortcut, nn.Conv2d): nn.init.kaiming_normal_(self.shortcut.weight, mode='fan_out', nonlinearity='leaky_relu') elif isinstance(self.shortcut, nn.ConvTranspose2d): nn.init.kaiming_normal_(self.shortcut.weight, mode='fan_out', nonlinearity='leaky_relu') if self.shortcut.bias is not None: nn.init.constant_(self.shortcut.bias, 0) def forward(self, x): identity = self.shortcut(x) # 主路径 x = self.conv_transpose(x) x = self.activation(x) x = self.conv1(x) x = self.activation(x) x = self.conv2(x) # 残差连接 x = x + identity x = self.activation(x) return x class FramePredictionDecoder(nn.Module): """Improved decoder for frame prediction""" def __init__(self, embed_dims, output_channels=1): super().__init__() # Reverse the embed_dims for decoder decoder_dims = embed_dims[::-1] self.blocks = nn.ModuleList() # 使用普通的DecoderBlock，第一个block使用大步长 self.blocks.append(DecoderBlock( decoder_dims[0], decoder_dims[1], kernel_size=3, stride=4, padding=1, output_padding=3 # 改为stride=4 )) self.blocks.append(DecoderBlock( decoder_dims[1], decoder_dims[2], kernel_size=3, stride=2, padding=1, output_padding=1 )) self.blocks.append(DecoderBlock( decoder_dims[2], decoder_dims[3], kernel_size=3, stride=2, padding=1, output_padding=1 )) # 第四个block：增加到64通道 self.blocks.append(DecoderBlock( decoder_dims[3], 64, kernel_size=3, stride=2, padding=1, output_padding=1 )) # 改进的最终输出层：不使用反卷积，只进行特征精炼 # 输入尺寸已经是目标尺寸，只需要调整通道数和进行特征融合 self.final_block = nn.Sequential( nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=True), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=True), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(32, output_channels, kernel_size=3, padding=1, bias=True) # 移除Tanh，让输出在任意范围，由损失函数和归一化处理 ) def forward(self, x): """ Args: x: input tensor of shape [B, embed_dims[-1], H/32, W/32] """ # 不使用skip connections for i in range(4): x = self.blocks[i](x) # 最终输出层：只进行特征精炼，不上采样 x = self.final_block(x) return x class SwiftFormerTemporal(nn.Module): """ SwiftFormer with temporal input for frame prediction. Input: [B, num_frames, H, W] (Y channel only) Output: predicted frame [B, 1, H, W] and optional representation """ def __init__(self, model_name='XS', num_frames=3, use_decoder=True, return_features=False, **kwargs): super().__init__() # Get model configuration layers = SwiftFormer_depth[model_name] embed_dims = SwiftFormer_width[model_name] # Store configuration self.num_frames = num_frames self.use_decoder = use_decoder self.return_features = return_features # Modify stem to accept multiple frames (only Y channel) in_channels = num_frames self.patch_embed = stem(in_channels, embed_dims[0]) # Build encoder network (same as SwiftFormer) network = [] for i in range(len(layers)): stage = Stage(embed_dims[i], i, layers, mlp_ratio=4, act_layer=nn.GELU, drop_rate=0., drop_path_rate=0., use_layer_scale=True, layer_scale_init_value=1e-5, vit_num=1) network.append(stage) if i >= len(layers) - 1: break if embed_dims[i] != embed_dims[i + 1]: network.append( Embedding( patch_size=3, stride=2, padding=1, in_chans=embed_dims[i], embed_dim=embed_dims[i + 1] ) ) self.network = nn.ModuleList(network) self.norm = nn.BatchNorm2d(embed_dims[-1]) # Frame prediction decoder if use_decoder: self.decoder = FramePredictionDecoder( embed_dims, output_channels=1 ) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, (nn.Conv2d, nn.Linear)): # 使用Kaiming初始化，适合ReLU/LeakyReLU nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.ConvTranspose2d): # 反卷积层使用特定的初始化 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='leaky_relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def forward_tokens(self, x): """Forward through encoder network, return list of stage features if return_features else final output""" if self.return_features: features = [] stage_idx = 0 for idx, block in enumerate(self.network): x = block(x) # 收集每个stage的输出（stage0, stage1, stage2, stage3） # 根据SwiftFormer结构，stage在索引0,2,4,6位置 if idx in [0, 2, 4, 6]: features.append(x) stage_idx += 1 return x, features else: for block in self.network: x = block(x) return x def forward(self, x): """ Args: x: input frames of shape [B, num_frames, H, W] Returns: If return_features is False: pred_frame: predicted frame [B, 1, H, W] (or None) If return_features is True: pred_frame, features (list of stage features) """ # Encode x = self.patch_embed(x) if self.return_features: x, features = self.forward_tokens(x) else: x = self.forward_tokens(x) x = self.norm(x) # Decode to frame pred_frame = None if self.use_decoder: pred_frame = self.decoder(x) if self.return_features: return pred_frame, features else: return pred_frame # Factory functions for different model sizes def SwiftFormerTemporal_XS(num_frames=3, **kwargs): return SwiftFormerTemporal('XS', num_frames=num_frames, **kwargs) def SwiftFormerTemporal_S(num_frames=3, **kwargs): return SwiftFormerTemporal('S', num_frames=num_frames, **kwargs) def SwiftFormerTemporal_L1(num_frames=3, **kwargs): return SwiftFormerTemporal('l1', num_frames=num_frames, **kwargs) def SwiftFormerTemporal_L3(num_frames=3, **kwargs): return SwiftFormerTemporal('l3', num_frames=num_frames, **kwargs)