1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164
| import torch import torch.nn as nn from typing import List, Tuple
class MViTBlock(nn.Module): """ Multiscale Vision Transformer Block MViT的核心:池化层次结构 - 逐层降低空间分辨率 - 逐层增加通道深度 - 捕获多尺度时空特征 """ def __init__( self, dim: int = 768, num_heads: int = 8, patch_size: Tuple[int, int, int] = (1, 14, 14), pooling_stride: Tuple[int, int, int] = (1, 2, 2), ): super().__init__() self.dim = dim self.num_heads = num_heads self.norm1 = nn.LayerNorm(dim) self.attn = nn.MultiheadAttention( embed_dim=dim, num_heads=num_heads, batch_first=True ) self.norm2 = nn.LayerNorm(dim) self.mlp = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), ) self.pool = nn.AvgPool3d( kernel_size=pooling_stride, stride=pooling_stride ) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Args: x: (B, T, H, W, C) 或 (B, N, C) token序列 Returns: pooled_features: (B, T', H', W', C') 降采样后 """ residual = x x = self.norm1(x) x_attn, _ = self.attn(x, x, x) x = residual + x_attn x = x + self.mlp(self.norm2(x)) if x.dim() == 5: x = x.permute(0, 4, 1, 2, 3) x = self.pool(x) x = x.permute(0, 2, 3, 4, 1) return x
class WeightSharedMViT(nn.Module): """ M2DAR核心架构:权重共享MViT处理多视角 同一个MViT模型处理3路视频流 → 学习视角不变的行为表示 → 减少参数量 → 防止对特定视角过拟合 """ def __init__(self, num_classes: int = 16, num_views: int = 3): super().__init__() self.num_views = num_views self.shared_backbone = nn.ModuleList([ MViTBlock(dim=768, num_heads=8), MViTBlock(dim=768, num_heads=8), MViTBlock(dim=768, num_heads=8), MViTBlock(dim=1536, num_heads=8), ]) self.cls_head = nn.Sequential( nn.Linear(1536, 768), nn.GELU(), nn.Dropout(0.5), nn.Linear(768, num_classes), nn.Sigmoid() ) self.tal_head = nn.Sequential( nn.Linear(1536, 512), nn.GELU(), nn.Linear(512, 2), ) def forward( self, views: List[torch.Tensor] ) -> Tuple[torch.Tensor, torch.Tensor]: """ Args: views: list of 3 tensors, each (B, C, T, H, W) Returns: cls_scores: (B, num_classes) tal_scores: (B, 2) start/end times """ view_features = [] for view in views: x = view B, C, T, H, W = x.shape x = x.reshape(B, C, T, H * W) x = x.permute(0, 2, 3, 1) for block in self.shared_backbone: x = block(x) feat = x.mean(dim=[1, 2]) view_features.append(feat) fused = torch.stack(view_features, dim=0).mean(dim=0) cls_scores = self.cls_head(fused) tal_scores = self.tal_head(fused) return cls_scores, tal_scores
if __name__ == "__main__": model = WeightSharedMViT(num_classes=16, num_views=3) B = 2 views = [ torch.randn(B, 3, 16, 224, 224) for _ in range(3) ] cls_scores, tal_scores = model(views) print(f"分类输出: {cls_scores.shape}") print(f"时序输出: {tal_scores.shape}") print(f"模型参数: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")
|