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
| import torch
import torch.nn as nn
import torch.nn.functional as F
class DiffAttnModel(nn.Module):
"""
DiffAttn: 基于扩散模型的驾驶员注意力预测
核心思想：
- 注意力图 S 服从高斯分布（类似人类注视模式）
- 扩散过程：S_t = sqrt(α_t)*S_0 + sqrt(1-α_t)*ε
- 去噪过程：从噪声注意力图逐步恢复清晰注意力分布
优势：
- 天然建模注视的高斯特性
- 概率生成式框架，非确定性回归
- 多尺度条件引导
"""
def __init__(self, num_timesteps=1000):
super().__init__()
self.num_timesteps = num_timesteps
self.encoder = SwinTransformerEncoder(
img_size=(384, 512),
embed_dim=128,
depths=[2, 2, 18, 2],
num_heads=[4, 8, 16, 32]
)
self.decoder = MultiScaleDiffusionDecoder(
in_channels=128,
mid_channels=64
)
self.llm_enhancer = LLMSemanticLayer(
feat_dim=128,
num_heads=8
)
self.noise_pred_head = nn.Conv2d(64, 1, 1)
def forward(self, image, t, noisy_attention):
"""
去噪一步
Args:
image: 驾驶场景图像 (B, 3, 384, 512)
t: 扩散时间步 (B,)
noisy_attention: 带噪声的注意力图 (B, 1, H, W)
Returns:
predicted_noise: 预测的噪声
"""
scene_features = self.encoder(image)
enhanced_features = self.llm_enhancer(scene_features)
denoised = self.decoder(enhanced_features, noisy_attention, t)
predicted_noise = self.noise_pred_head(denoised)
return predicted_noise
def generate_attention(self, image):
"""
生成驾驶员注意力图（推理时逐步去噪）
"""
B, _, H, W = image.shape
attention = torch.randn(B, 1, H, W).to(image.device)
for t in reversed(range(self.num_timesteps)):
t_tensor = torch.full((B,), t, device=image.device)
predicted_noise = self.forward(image, t_tensor, attention)
attention = self.ddpm_step(attention, predicted_noise, t)
return attention.clamp(0, 1)
def ddpm_step(self, x, predicted_noise, t):
"""DDPM采样一步"""
beta = self.betas[t]
alpha = self.alphas[t]
alpha_bar = self.alpha_bars[t]
mean = (1 / alpha.sqrt()) * (x - beta / (1 - alpha_bar).sqrt() * predicted_noise)
if t > 0:
noise = torch.randn_like(x)
mean += beta.sqrt() * noise
return mean
class MultiScaleDiffusionDecoder(nn.Module):
"""多尺度条件去噪解码器"""
def __init__(self, in_channels, mid_channels):
super().__init__()
self.ffp = FeatureFusionPyramid(in_channels, mid_channels)
self.condition_inject = nn.ModuleList([
nn.Conv2d(mid_channels, mid_channels, 3, 1, 1)
for _ in range(4)
])
def forward(self, features, noisy_attn, t):
fused = self.ffp(features)
return torch.cat([fused, noisy_attn], dim=1)
class LLMSemanticLayer(nn.Module):
"""LLM增强语义推理层"""
def __init__(self, feat_dim, num_heads):
super().__init__()
self.cross_attn = nn.MultiheadAttention(
feat_dim, num_heads, batch_first=True
)
self.norm = nn.LayerNorm(feat_dim)
self.ffn = nn.Sequential(
nn.Linear(feat_dim, feat_dim * 2),
nn.GELU(),
nn.Linear(feat_dim * 2, feat_dim)
)
def forward(self, visual_features):
"""
用LLM语义知识增强视觉特征
论文核心：不显式引入分割图/光流等top-down特征
而是通过LLM层隐式增强语义推理
"""
B, C, H, W = visual_features.shape
x = visual_features.flatten(2).permute(0, 2, 1)
attn_out, _ = self.cross_attn(x, x, x)
x = self.norm(x + attn_out)
x = self.norm(x + self.ffn(x))
return x.permute(0, 2, 1).reshape(B, C, H, W)
|
