""" 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 DecoderBlockWithSkip(nn.Module): """Decoder block with skip connection support""" def __init__(self, in_channels, out_channels, skip_channels=0, kernel_size=3, stride=2, padding=1, output_padding=1): super().__init__() # 总输入通道 = 输入通道 + skip通道 total_in_channels = in_channels + skip_channels # 主路径：反卷积 + 两个卷积层 self.conv_transpose = nn.ConvTranspose2d( total_in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, output_padding=output_padding, bias=True ) 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 total_in_channels != out_channels or stride != 1: if stride == 1: self.shortcut = nn.Conv2d(total_in_channels, out_channels, kernel_size=1, bias=True) else: self.shortcut = nn.ConvTranspose2d( total_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, skip_feature=None): # 如果有skip feature，将其与输入拼接 if skip_feature is not None: # 确保skip特征的空间尺寸与x匹配 if skip_feature.shape[2:] != x.shape[2:]: # 使用双线性插值进行上采样或下采样 skip_feature = torch.nn.functional.interpolate( skip_feature, size=x.shape[2:], mode='bilinear', align_corners=False ) x = torch.cat([x, skip_feature], dim=1) 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 with better upsampling strategy""" def __init__(self, embed_dims, output_channels=1, use_skip=False): super().__init__() self.use_skip = use_skip # Reverse the embed_dims for decoder decoder_dims = embed_dims[::-1] self.blocks = nn.ModuleList() if use_skip: # 使用支持skip connections的block # 第一个block：从bottleneck到stage4，使用大步长stride=4，skip来自stage3 self.blocks.append(DecoderBlockWithSkip( decoder_dims[0], decoder_dims[1], skip_channels=embed_dims[3], # stage3的通道数 kernel_size=3, stride=4, padding=1, output_padding=3 # 改为stride=4 )) # 第二个block：stage4到stage3，stride=2，skip来自stage2 self.blocks.append(DecoderBlockWithSkip( decoder_dims[1], decoder_dims[2], skip_channels=embed_dims[2], # stage2的通道数 kernel_size=3, stride=2, padding=1, output_padding=1 )) # 第三个block：stage3到stage2，stride=2，skip来自stage1 self.blocks.append(DecoderBlockWithSkip( decoder_dims[2], decoder_dims[3], skip_channels=embed_dims[1], # stage1的通道数 kernel_size=3, stride=2, padding=1, output_padding=1 )) # 第四个block：stage2到stage1，stride=2，skip来自stage0 self.blocks.append(DecoderBlockWithSkip( decoder_dims[3], 64, # 输出到64通道 skip_channels=embed_dims[0], # stage0的通道数 kernel_size=3, stride=2, padding=1, output_padding=1 )) else: # 使用普通的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, skip_features=None): """ Args: x: input tensor of shape [B, embed_dims[-1], H/32, W/32] skip_features: list of encoder features from stages [stage3, stage2, stage1, stage0] each of shape [B, C, H', W'] where C matches encoder dims """ if self.use_skip: if skip_features is None: raise ValueError("skip_features must be provided when use_skip=True") # 确保有4个skip features assert len(skip_features) == 4, f"Need 4 skip features, got {len(skip_features)}" # 反转顺序以匹配解码器：stage3, stage2, stage1, stage0 skip_features = skip_features[::-1] # 调整skip特征的尺寸以匹配新的上采样策略 adjusted_skip_features = [] for i, skip in enumerate(skip_features): if skip is not None: # 计算目标尺寸：4, 2, 2, 2倍上采样 upsample_factors = [4, 2, 2, 2] target_height = x.shape[2] * upsample_factors[i] target_width = x.shape[3] * upsample_factors[i] if skip.shape[2:] != (target_height, target_width): skip = torch.nn.functional.interpolate( skip, size=(target_height, target_width), mode='bilinear', align_corners=False ) adjusted_skip_features.append(skip) # 四个block使用skip connections for i in range(4): x = self.blocks[i](x, adjusted_skip_features[i]) else: # 不使用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, use_skip=True, # 新增：是否使用skip connections use_representation_head=False, representation_dim=128, 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.use_skip = use_skip # 保存skip connections设置 self.use_representation_head = use_representation_head 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, use_skip=use_skip # 传递skip connections设置 ) # Representation head for pose/velocity prediction if use_representation_head: self.representation_head = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(embed_dims[-1], representation_dim), nn.ReLU(), nn.Linear(representation_dim, representation_dim) ) else: self.representation_head = None 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 or self.use_skip: 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) representation: optional representation vector [B, representation_dim] (or None) If return_features is True: pred_frame, representation, features (list of stage features) """ # Encode x = self.patch_embed(x) if self.return_features or self.use_skip: x, features = self.forward_tokens(x) else: x = self.forward_tokens(x) x = self.norm(x) # Get representation if needed representation = None if self.representation_head is not None: representation = self.representation_head(x) # Decode to frame pred_frame = None if self.use_decoder: if self.use_skip: # 提取用于skip connections的特征 # features包含所有stage的输出，我们需要stage0, stage1, stage2, stage3 # 根据SwiftFormer结构，应该有4个stage特征 if len(features) >= 4: # 取四个stage的特征：stage0, stage1, stage2, stage3 skip_features = [features[0], features[1], features[2], features[3]] else: # 如果特征不够，使用可用的特征 skip_features = features[:4] # 如果特征仍然不够，使用None填充 while len(skip_features) < 4: skip_features.append(None) pred_frame = self.decoder(x, skip_features) else: pred_frame = self.decoder(x) if self.return_features: return pred_frame, representation, features else: return pred_frame, representation # Factory functions for different model sizes def SwiftFormerTemporal_XS(num_frames=3, use_skip=True, **kwargs): return SwiftFormerTemporal('XS', num_frames=num_frames, use_skip=use_skip, **kwargs) def SwiftFormerTemporal_S(num_frames=3, use_skip=True, **kwargs): return SwiftFormerTemporal('S', num_frames=num_frames, use_skip=use_skip, **kwargs) def SwiftFormerTemporal_L1(num_frames=3, use_skip=True, **kwargs): return SwiftFormerTemporal('l1', num_frames=num_frames, use_skip=use_skip, **kwargs) def SwiftFormerTemporal_L3(num_frames=3, use_skip=True, **kwargs): return SwiftFormerTemporal('l3', num_frames=num_frames, use_skip=use_skip, **kwargs)