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
| class GDSQ(nn.Module): """ Gaze-Driven Semantic Query Module 核心创新:用眼动引导视频token选择 思路: - 驾驶员看哪,就从视频中选择对应的视觉token - 实现眼动-场景上下文交互建模 """ def __init__(self, video_dim=768, gaze_dim=256, num_heads=8): super().__init__() self.cross_attention = nn.MultiheadAttention( embed_dim=video_dim, num_heads=num_heads, batch_first=True ) self.gaze_proj = nn.Linear(gaze_dim, video_dim) self.gate = nn.Sequential( nn.Linear(video_dim * 2, video_dim), nn.Sigmoid() ) def forward(self, video_tokens, gaze_features, gaze_positions): """ 前向传播 Args: video_tokens: (B, N, 768) 视频token序列 gaze_features: (B, 256) 眼动特征 gaze_positions: (B, T, 2) 注视点位置(归一化) Returns: enhanced_features: (B, 768) 注视增强特征 """ B, N, D = video_tokens.shape gaze_query = self.gaze_proj(gaze_features).unsqueeze(1) attn_out, attn_weights = self.cross_attention( query=gaze_query, key=video_tokens, value=video_tokens ) global_feature = video_tokens.mean(dim=1) combined = torch.cat([attn_out.squeeze(1), global_feature], dim=-1) gate_weights = self.gate(combined) enhanced_features = gate_weights * attn_out.squeeze(1) + (1 - gate_weights) * global_feature return enhanced_features, attn_weights
gdsq = GDSQ() video_tokens = torch.randn(2, 196, 768) gaze_features = torch.randn(2, 256) gaze_positions = torch.rand(2, 100, 2)
enhanced_features, attn_weights = gdsq(video_tokens, gaze_features, gaze_positions) print(f"增强特征: {enhanced_features.shape}, 注意力权重: {attn_weights.shape}")
|