Star845
1320
代码
代码
Issues1488
Pull Requests2327
流水线
项目讨论
Wiki
分析
Star845
1320
已合并
update: ads model #38091
bigprestigee1创建于 6月10日
update: ads model #38091
已合并
bigprestigee1创建于 6月10日
从bigprestigee1/pytorch_commonv2.7.1_ads合入到Ascend/pytorchv2.7.1
共 8 个文件变更+1510-0
Abenchmarks/v1sparsedrive/README.md+71-0
@@ -0,0 +1,71 @@
1+# SimpleV1SparseDrive (torch-npu / Ascend)
2+ 
3+---
4+ 
5+## 1. 项目概述
6+ 
7+本项目实现了名为 **SimpleV1SparseDrive** 的模型结构，整体由以下模块组成：
8+ 
9+| 模块 | 说明 |
10+|------|------|
11+| `data_preprocessor` | `BaseDataPreprocessor()` |
12+| `img_backbone` | ResNet（包含大量 `BatchNorm2d`） |
13+| `img_neck` | FPN |
14+| `head` | `V1SparseDriveHead` |
15+| ├─ `det_head` | Sparse4DHeadLike（含 DeformableFeatureAggregation / FlashAttention / Refinement） |
16+| ├─ `map_head` | Sparse4DHeadLike |
17+| └─ `motion_plan_head` | Motion & Planning（含多层 FlashAttention / FFN / refinement） |
18+| `depth_branch` | DenseDepthNetLike（3 个 depth layer） |
19+| `grid_mask` | GridMask |
20+ 
21+---
22+ 
23+## 2. 环境与依赖
24+ 
25+### 2.1 硬件/系统要求
26+ 
27+- Ascend NPU 环境已正确安装（驱动 / CANN 等）
28+ 
29+### 2.2 Pip 环境包（已验证可复现）
30+ 
31+- 基础依赖包
32+ 
33+```bash
34+pip install -r requirement.txt
35+```
36+ 
37+- 源码安装mmcv
38+ 
39+```bash
40+git clone -b 1.x https://github.com/open-mmlab/mmcv.git
41+cd mmcv
42+MMCV_WITH_OPS=1 FORCE_NPU=1 python setup.py install
43+cd ..
44+```
45+ 
46+> 建议固定上述版本以保证结果可复现。
47+ 
48+---
49+ 
50+## 3. 运行训练和验证
51+ 
52+本工程提供两种运行模式：
53+ 
54+- **Eager（默认）**：不启用 `torch.compile`，直接动态图执行。
55+- **Graph/Compile**：通过 `--compile` 启用 `torch.compile`（Inductor 图模式），首步通常包含编译开销。
56+ 
57+> 默认运行 **200 step**（如需修改步数，请按工程内 `config.py` 的参数/配置为准）。
58+ 
59+### 3.1 训练（生成日志）
60+ 
61+#### 3.1.1 Eager 模式（默认）
62+ 
63+```bash
64+python train.py 2>&1 | tee eager.log
65+```
66+ 
67+#### 3.1.2 Graph / torch.compile 模式
68+ 
69+```bash
70+python train.py --npu-backend dvm --compile 2>&1 | tee compile.log
71+```
Abenchmarks/v1sparsedrive/config.py+25-0
@@ -0,0 +1,25 @@
1+CONFIG = dict(
2+ seed=3407,
3+ batch_size=2,
4+ num_workers=0,
5+ max_epochs=1,
6+ train_steps=200,
7+ lr=1e-4,
8+ weight_decay=1e-4,
9+ 
10+ model=dict(
11+ num_det_classes=10,
12+ num_map_classes=3,
13+ backbone_pretrained=None,
14+ use_grid_mask=True,
15+ ),
16+ 
17+ dummy_dataset=dict(
18+ length=400,
19+ image_size=(256, 704),
20+ num_det_boxes=20,
21+ num_map_instances=10,
22+ num_det_classes=10,
23+ num_map_classes=3,
24+ )
25+)
Abenchmarks/v1sparsedrive/datasets/__init__.py+3-0
@@ -0,0 +1,3 @@
1+from .dummy_dataset import DummySparseDriveDataset, simple_collate
2+ 
3+__all__ = ["DummySparseDriveDataset", "simple_collate"]
Abenchmarks/v1sparsedrive/datasets/dummy_dataset.py+77-0
@@ -0,0 +1,77 @@
1+import torch
2+from torch.utils.data import Dataset
3+ 
4+ 
5+class DummySparseDriveDataset(Dataset):
6+ """
7+ 模拟后续真实 dataloader 输出格式：
8+ {
9+ "inputs": Tensor[3, H, W],
10+ "data_samples": {
11+ "sample_idx": int,
12+ "img_shape": (H, W),
13+ "gt_det_boxes": Tensor[num_det_boxes, 11],
14+ "gt_det_labels": Tensor[num_det_boxes],
15+ "gt_map_points": Tensor[num_map_instances, 40],
16+ "gt_map_labels": Tensor[num_map_instances],
17+ "gt_motion": Tensor[24],
18+ "gt_plan": Tensor[12],
19+ }
20+ }
21+ """
22+ def __init__(
23+ self,
24+ length=20,
25+ image_size=(256, 704),
26+ num_det_boxes=20,
27+ num_map_instances=10,
28+ num_det_classes=10,
29+ num_map_classes=3,
30+ ):
31+ self.length = length
32+ self.image_size = image_size
33+ self.num_det_boxes = num_det_boxes
34+ self.num_map_instances = num_map_instances
35+ self.num_det_classes = num_det_classes
36+ self.num_map_classes = num_map_classes
37+ 
38+ def __len__(self):
39+ return self.length
40+ 
41+ def __getitem__(self, idx):
42+ h, w = self.image_size
43+ 
44+ image = torch.randn(3, h, w)
45+ 
46+ gt_det_boxes = torch.randn(self.num_det_boxes, 11)
47+ gt_det_labels = torch.randint(0, self.num_det_classes, (self.num_det_boxes,))
48+ 
49+ gt_map_points = torch.randn(self.num_map_instances, 40)
50+ gt_map_labels = torch.randint(0, self.num_map_classes, (self.num_map_instances,))
51+ 
52+ gt_motion = torch.randn(24)
53+ gt_plan = torch.randn(12)
54+ 
55+ sample = {
56+ "inputs": image,
57+ "data_samples": {
58+ "sample_idx": idx,
59+ "img_shape": (h, w),
60+ "gt_det_boxes": gt_det_boxes,
61+ "gt_det_labels": gt_det_labels,
62+ "gt_map_points": gt_map_points,
63+ "gt_map_labels": gt_map_labels,
64+ "gt_motion": gt_motion,
65+ "gt_plan": gt_plan,
66+ }
67+ }
68+ return sample
69+ 
70+ 
71+def simple_collate(batch):
72+ inputs = torch.stack([item["inputs"] for item in batch], dim=0)
73+ data_samples = [item["data_samples"] for item in batch]
74+ return {
75+ "inputs": inputs,
76+ "data_samples": data_samples,
77+ }
Abenchmarks/v1sparsedrive/models/__init__.py+3-0
@@ -0,0 +1,3 @@
1+from .simple_v1sparsedrive import SimpleV1SparseDrive
2+ 
3+__all__ = ["SimpleV1SparseDrive"]
Abenchmarks/v1sparsedrive/models/simple_v1sparsedrive.py+1045-0
@@ -0,0 +1,1045 @@
1+import math
2+import torch
3+import torch.nn as nn
4+import torch.nn.functional as F
5+ 
6+from mmengine.model import BaseModel
7+from mmdet.models.backbones import ResNet
8+from mmdet.models.necks import FPN
9+ 
10+ 
11+# =========================
12+# utils
13+# =========================
14+ 
15+def build_mlp(channels, act_layer=nn.ReLU, last_norm=False):
16+ layers = []
17+ for i in range(len(channels) - 1):
18+ in_c = channels[i]
19+ out_c = channels[i + 1]
20+ layers.append(nn.Linear(in_c, out_c))
21+ is_last = (i == len(channels) - 2)
22+ if not is_last:
23+ layers.append(act_layer(inplace=True) if act_layer == nn.ReLU else act_layer())
24+ layers.append(nn.LayerNorm(out_c))
25+ elif last_norm:
26+ layers.append(nn.LayerNorm(out_c))
27+ return nn.Sequential(*layers)
28+ 
29+ 
30+class Scale(nn.Module):
31+ def __init__(self, init_value=1.0):
32+ super().__init__()
33+ self.scale = nn.Parameter(torch.tensor(float(init_value)))
34+ 
35+ def forward(self, x):
36+ return x * self.scale
37+ 
38+ 
39+# =========================
40+# grid mask
41+# =========================
42+ 
43+class GridMask(nn.Module):
44+ def __init__(self, use_h=True, use_w=True, rotate=1, ratio=0.5, prob=0.7):
45+ super().__init__()
46+ self.use_h = use_h
47+ self.use_w = use_w
48+ self.rotate = rotate
49+ self.ratio = ratio
50+ self.prob = prob
51+ 
52+ def forward(self, x):
53+ if (not self.training) or torch.rand(1).item() > self.prob:
54+ return x
55+ 
56+ _, _, h, w = x.shape
57+ device = x.device
58+ 
59+ d = max(2, min(h, w) // 8)
60+ l = max(1, int(d * self.ratio))
61+ 
62+ mask = torch.ones((h, w), device=device, dtype=x.dtype)
63+ 
64+ if self.use_h:
65+ for i in range(0, h, d):
66+ mask[i:i + l, :] = 0
67+ 
68+ if self.use_w:
69+ for j in range(0, w, d):
70+ mask[:, j:j + l] = 0
71+ 
72+ mask = mask.unsqueeze(0).unsqueeze(0)
73+ x = x * mask
74+ return x
75+ 
76+ 
77+# =========================
78+# anchor / encoder like
79+# =========================
80+ 
81+class SparseBox3DKeyPointsGenerator(nn.Module):
82+ def __init__(self, embed_dims=256, num_pts=9):
83+ super().__init__()
84+ self.learnable_fc = nn.Linear(embed_dims, num_pts * 2)
85+ 
86+ def forward(self, query):
87+ return self.learnable_fc(query)
88+ 
89+ 
90+class SparsePoint3DKeyPointsGenerator(nn.Module):
91+ def __init__(self, embed_dims=256, out_dim=600):
92+ super().__init__()
93+ self.learnable_fc = nn.Linear(embed_dims, out_dim)
94+ 
95+ def forward(self, query):
96+ return self.learnable_fc(query)
97+ 
98+ 
99+class TrajSparsePoint3DKeyPointsGenerator(nn.Module):
100+ def __init__(self):
101+ super().__init__()
102+ 
103+ def forward(self, x):
104+ return x
105+ 
106+ 
107+class InstanceBank(nn.Module):
108+ def __init__(self, num_queries=300, embed_dims=256, anchor_handler=None):
109+ super().__init__()
110+ self.num_queries = num_queries
111+ self.embed_dims = embed_dims
112+ self.anchor_handler = anchor_handler
113+ self.query_embed = nn.Parameter(torch.randn(num_queries, embed_dims) * 0.02)
114+ 
115+ def forward(self, batch_size, device):
116+ query = self.query_embed.unsqueeze(0).repeat(batch_size, 1, 1).to(device)
117+ return query
118+ 
119+ 
120+class SparseBox3DEncoder(nn.Module):
121+ def __init__(self):
122+ super().__init__()
123+ self.pos_fc = nn.Sequential(
124+ nn.Linear(3, 128), nn.ReLU(inplace=True), nn.LayerNorm(128),
125+ nn.Linear(128, 128), nn.ReLU(inplace=True), nn.LayerNorm(128),
126+ nn.Linear(128, 128), nn.ReLU(inplace=True), nn.LayerNorm(128),
127+ nn.Linear(128, 128), nn.ReLU(inplace=True), nn.LayerNorm(128),
128+ )
129+ self.size_fc = nn.Sequential(
130+ nn.Linear(3, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
131+ nn.Linear(32, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
132+ nn.Linear(32, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
133+ nn.Linear(32, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
134+ )
135+ self.yaw_fc = nn.Sequential(
136+ nn.Linear(2, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
137+ nn.Linear(32, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
138+ nn.Linear(32, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
139+ nn.Linear(32, 32), nn.ReLU(inplace=True), nn.LayerNorm(32),
140+ )
141+ self.vel_fc = nn.Sequential(
142+ nn.Linear(3, 64), nn.ReLU(inplace=True), nn.LayerNorm(64),
143+ nn.Linear(64, 64), nn.ReLU(inplace=True), nn.LayerNorm(64),
144+ nn.Linear(64, 64), nn.ReLU(inplace=True), nn.LayerNorm(64),
145+ nn.Linear(64, 64), nn.ReLU(inplace=True), nn.LayerNorm(64),
146+ )
147+ self.out_proj = nn.Linear(128 + 32 + 32 + 64, 256)
148+ 
149+ def forward(self, anchors):
150+ pos = anchors[..., 0:3]
151+ size = anchors[..., 3:6]
152+ yaw = anchors[..., 6:8]
153+ vel = anchors[..., 8:11]
154+ x = torch.cat([
155+ self.pos_fc(pos),
156+ self.size_fc(size),
157+ self.yaw_fc(yaw),
158+ self.vel_fc(vel),
159+ ], dim=-1)
160+ return self.out_proj(x)
161+ 
162+ 
163+class SparsePoint3DEncoder(nn.Module):
164+ def __init__(self):
165+ super().__init__()
166+ self.pos_fc = nn.Sequential(
167+ nn.Linear(40, 256),
168+ nn.ReLU(inplace=True),
169+ nn.LayerNorm(256),
170+ nn.Linear(256, 256),
171+ nn.ReLU(inplace=True),
172+ nn.LayerNorm(256),
173+ )
174+ 
175+ def forward(self, anchors):
176+ return self.pos_fc(anchors)
177+ 
178+ 
179+# =========================
180+# attention / ffn like
181+# =========================
182+ 
183+class MultiheadFlashAttention(nn.Module):
184+ def __init__(self, embed_dims=256, num_heads=8, batch_first=True):
185+ super().__init__()
186+ self.attn = nn.MultiheadAttention(
187+ embed_dim=embed_dims,
188+ num_heads=num_heads,
189+ dropout=0.1,
190+ batch_first=batch_first,
191+ )
192+ self.proj_drop = nn.Dropout(0.0)
193+ self.dropout_layer = nn.Dropout(0.1)
194+ 
195+ def forward(self, x):
196+ out, _ = self.attn(x, x, x, need_weights=False)
197+ out = self.proj_drop(out)
198+ out = self.dropout_layer(out)
199+ return x + out
200+ 
201+ 
202+class MultiheadAttention(nn.Module):
203+ def __init__(self, embed_dims=256, num_heads=8, kdim=None, vdim=None, batch_first=True):
204+ super().__init__()
205+ self.attn = nn.MultiheadAttention(
206+ embed_dim=embed_dims,
207+ num_heads=num_heads,
208+ dropout=0.1,
209+ batch_first=batch_first,
210+ kdim=kdim,
211+ vdim=vdim,
212+ )
213+ self.proj_drop = nn.Dropout(0.0)
214+ self.dropout_layer = nn.Dropout(0.1)
215+ 
216+ def forward(self, q, k=None, v=None):
217+ if k is None:
218+ k = q
219+ if v is None:
220+ v = k
221+ out, _ = self.attn(q, k, v, need_weights=False)
222+ out = self.proj_drop(out)
223+ out = self.dropout_layer(out)
224+ return q + out
225+ 
226+ 
227+class AsymmetricFFN(nn.Module):
228+ def __init__(self, in_dims=256, embed_dims=256, hidden_dims=512, pre_norm_dim=None):
229+ super().__init__()
230+ if pre_norm_dim is None:
231+ pre_norm_dim = in_dims
232+ 
233+ self.in_dims = in_dims
234+ self.embed_dims = embed_dims
235+ self.pre_norm_dim = pre_norm_dim
236+ 
237+ self.pre_norm = nn.LayerNorm(pre_norm_dim)
238+ self.layers = nn.Sequential(
239+ nn.Sequential(
240+ nn.Linear(in_dims, hidden_dims),
241+ nn.ReLU(inplace=True),
242+ nn.Dropout(0.1),
243+ ),
244+ nn.Linear(hidden_dims, embed_dims),
245+ nn.Dropout(0.1),
246+ )
247+ self.dropout_layer = nn.Identity()
248+ self.identity_fc = nn.Linear(in_dims, embed_dims) if in_dims != embed_dims else nn.Identity()
249+ 
250+ def forward(self, x):
251+ identity = self.identity_fc(x)
252+ out = self.pre_norm(x)
253+ out = self.layers(out)
254+ out = self.dropout_layer(out)
255+ return identity + out
256+ 
257+ 
258+# =========================
259+# feature aggregation like
260+# =========================
261+ 
262+class DeformableFeatureAggregation(nn.Module):
263+ def __init__(self, embed_dims=256, num_levels=4, num_pts_out=18, weights_out=416):
264+ super().__init__()
265+ self.proj_drop = nn.Dropout(0.0)
266+ self.kps_generator = SparseBox3DKeyPointsGenerator(embed_dims, num_pts=num_pts_out // 2) \
267+ if num_pts_out == 18 else SparsePoint3DKeyPointsGenerator(embed_dims, out_dim=num_pts_out)
268+ self.output_proj = nn.Linear(embed_dims, embed_dims)
269+ self.camera_encoder = nn.Sequential(
270+ nn.Linear(12, 256),
271+ nn.ReLU(inplace=True),
272+ nn.LayerNorm(256),
273+ nn.Linear(256, 256),
274+ nn.ReLU(inplace=True),
275+ nn.LayerNorm(256),
276+ )
277+ self.weights_fc = nn.Linear(embed_dims, weights_out)
278+ 
279+ self.level_proj = nn.ModuleList([
280+ nn.Linear(embed_dims, embed_dims) for _ in range(num_levels)
281+ ])
282+ 
283+ def forward(self, query, fpn_feats, metas=None):
284+ pooled = []
285+ for i, feat in enumerate(fpn_feats):
286+ x = F.adaptive_avg_pool2d(feat, (1, 1)).flatten(1)
287+ x = self.level_proj[i](x).unsqueeze(1)
288+ pooled.append(x)
289+ 
290+ pooled = torch.stack(pooled, dim=0).mean(dim=0)
291+ query = query + pooled
292+ 
293+ out = self.output_proj(query)
294+ out = self.proj_drop(out)
295+ return out
296+ 
297+ 
298+# =========================
299+# refinement like
300+# =========================
301+ 
302+class SparseBox3DRefinementModule(nn.Module):
303+ def __init__(self, embed_dims=256, num_classes=10, box_dim=11):
304+ super().__init__()
305+ self.layers = nn.Sequential(
306+ nn.Linear(embed_dims, embed_dims),
307+ nn.ReLU(inplace=True),
308+ nn.Linear(embed_dims, embed_dims),
309+ nn.ReLU(inplace=True),
310+ nn.LayerNorm(embed_dims),
311+ nn.Linear(embed_dims, embed_dims),
312+ nn.ReLU(inplace=True),
313+ nn.Linear(embed_dims, embed_dims),
314+ nn.ReLU(inplace=True),
315+ nn.LayerNorm(embed_dims),
316+ nn.Linear(embed_dims, box_dim),
317+ Scale(),
318+ )
319+ self.cls_layers = nn.Sequential(
320+ nn.Linear(embed_dims, embed_dims),
321+ nn.ReLU(inplace=True),
322+ nn.LayerNorm(embed_dims),
323+ nn.Linear(embed_dims, embed_dims),
324+ nn.ReLU(inplace=True),
325+ nn.LayerNorm(embed_dims),
326+ nn.Linear(embed_dims, num_classes),
327+ )
328+ self.quality_layers = nn.Sequential(
329+ nn.Linear(embed_dims, embed_dims),
330+ nn.ReLU(inplace=True),
331+ nn.LayerNorm(embed_dims),
332+ nn.Linear(embed_dims, embed_dims),
333+ nn.ReLU(inplace=True),
334+ nn.LayerNorm(embed_dims),
335+ nn.Linear(embed_dims, 2),
336+ )
337+ 
338+ def forward(self, query):
339+ box = self.layers(query)
340+ cls = self.cls_layers(query)
341+ quality = self.quality_layers(query)
342+ return box, cls, quality
343+ 
344+ 
345+class SparsePoint3DRefinementModule(nn.Module):
346+ def __init__(self, embed_dims=256, num_classes=3, point_dim=40):
347+ super().__init__()
348+ self.layers = nn.Sequential(
349+ nn.Linear(embed_dims, embed_dims),
350+ nn.ReLU(inplace=True),
351+ nn.Linear(embed_dims, embed_dims),
352+ nn.ReLU(inplace=True),
353+ nn.LayerNorm(embed_dims),
354+ nn.Linear(embed_dims, embed_dims),
355+ nn.ReLU(inplace=True),
356+ nn.Linear(embed_dims, embed_dims),
357+ nn.ReLU(inplace=True),
358+ nn.LayerNorm(embed_dims),
359+ nn.Linear(embed_dims, point_dim),
360+ Scale(),
361+ )
362+ self.cls_layers = nn.Sequential(
363+ nn.Linear(embed_dims, embed_dims),
364+ nn.ReLU(inplace=True),
365+ nn.LayerNorm(embed_dims),
366+ nn.Linear(embed_dims, embed_dims),
367+ nn.ReLU(inplace=True),
368+ nn.LayerNorm(embed_dims),
369+ nn.Linear(embed_dims, num_classes),
370+ )
371+ 
372+ def forward(self, query):
373+ pts = self.layers(query)
374+ cls = self.cls_layers(query)
375+ return pts, cls
376+ 
377+ 
378+# =========================
379+# sparse4d head like
380+# =========================
381+ 
382+class Sparse4DHeadLike(nn.Module):
383+ def __init__(
384+ self,
385+ mode="det",
386+ embed_dims=256,
387+ num_queries=300,
388+ num_classes=10,
389+ num_decoder=6,
390+ ):
391+ super().__init__()
392+ self.mode = mode
393+ 
394+ if mode == "det":
395+ self.instance_bank = InstanceBank(
396+ num_queries=num_queries,
397+ embed_dims=embed_dims,
398+ anchor_handler=SparseBox3DKeyPointsGenerator(),
399+ )
400+ self.anchor_encoder = SparseBox3DEncoder()
401+ self.fc_before = nn.Linear(embed_dims, 512, bias=False)
402+ self.fc_after = nn.Linear(512, embed_dims, bias=False)
403+ 
404+ layers = []
405+ for _ in range(num_decoder):
406+ layers.extend([
407+ DeformableFeatureAggregation(embed_dims=256, num_pts_out=18, weights_out=416),
408+ AsymmetricFFN(in_dims=512, embed_dims=256, hidden_dims=1024, pre_norm_dim=512),
409+ nn.LayerNorm(256),
410+ SparseBox3DRefinementModule(embed_dims=256, num_classes=num_classes, box_dim=11),
411+ MultiheadFlashAttention(embed_dims=512, num_heads=8),
412+ MultiheadFlashAttention(embed_dims=512, num_heads=8),
413+ nn.LayerNorm(256),
414+ ])
415+ self.layers = nn.ModuleList(layers)
416+ 
417+ else:
418+ self.instance_bank = InstanceBank(
419+ num_queries=num_queries,
420+ embed_dims=embed_dims,
421+ anchor_handler=SparsePoint3DKeyPointsGenerator(),
422+ )
423+ self.anchor_encoder = SparsePoint3DEncoder()
424+ self.fc_before = nn.Identity()
425+ self.fc_after = nn.Identity()
426+ 
427+ layers = []
428+ layers.extend([
429+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
430+ nn.LayerNorm(256),
431+ ])
432+ for _ in range(num_decoder):
433+ layers.extend([
434+ DeformableFeatureAggregation(embed_dims=256, num_pts_out=600, weights_out=9600),
435+ AsymmetricFFN(in_dims=256, embed_dims=256, hidden_dims=512, pre_norm_dim=256),
436+ nn.LayerNorm(256),
437+ SparsePoint3DRefinementModule(embed_dims=256, num_classes=num_classes, point_dim=40),
438+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
439+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
440+ nn.LayerNorm(256),
441+ ])
442+ self.layers = nn.ModuleList(layers)
443+ 
444+ def _make_dummy_anchors(self, batch_size, num_queries, device):
445+ if self.mode == "det":
446+ return torch.randn(batch_size, num_queries, 11, device=device)
447+ else:
448+ return torch.randn(batch_size, num_queries, 40, device=device)
449+ 
450+ def forward(self, fpn_feats, metas=None):
451+ b = fpn_feats[0].shape[0]
452+ device = fpn_feats[0].device
453+ query = self.instance_bank(b, device)
454+ 
455+ anchors = self._make_dummy_anchors(b, query.shape[1], device)
456+ anchor_embed = self.anchor_encoder(anchors)
457+ query = query + anchor_embed
458+ 
459+ all_cls = []
460+ all_reg = []
461+ all_quality = []
462+ 
463+ i = 0
464+ while i < len(self.layers):
465+ layer = self.layers[i]
466+ 
467+ if isinstance(layer, DeformableFeatureAggregation):
468+ q_in = self.fc_before(query)
469+ q = self.fc_after(q_in)
470+ query = layer(q, fpn_feats, metas)
471+ i += 1
472+ 
473+ elif isinstance(layer, AsymmetricFFN):
474+ q_in = self.fc_before(query)
475+ query = layer(q_in)
476+ i += 1
477+ 
478+ elif isinstance(layer, nn.LayerNorm):
479+ query = layer(query)
480+ i += 1
481+ 
482+ elif isinstance(layer, SparseBox3DRefinementModule):
483+ box, cls, quality = layer(query)
484+ all_reg.append(box)
485+ all_cls.append(cls)
486+ all_quality.append(quality)
487+ i += 1
488+ 
489+ elif isinstance(layer, SparsePoint3DRefinementModule):
490+ pts, cls = layer(query)
491+ all_reg.append(pts)
492+ all_cls.append(cls)
493+ i += 1
494+ 
495+ elif isinstance(layer, MultiheadFlashAttention):
496+ q_in = self.fc_before(query)
497+ q_out = layer(q_in)
498+ query = self.fc_after(q_out)
499+ i += 1
500+ 
501+ else:
502+ query = layer(query)
503+ i += 1
504+ 
505+ out = {
506+ "query": query,
507+ "all_cls_scores": all_cls,
508+ "all_reg_preds": all_reg,
509+ }
510+ if self.mode == "det":
511+ out["all_quality_scores"] = all_quality
512+ return out
513+ 
514+ 
515+# =========================
516+# motion plan head like
517+# =========================
518+ 
519+class InstanceQueue(nn.Module):
520+ def __init__(self, embed_dims=256):
521+ super().__init__()
522+ self.ego_feature_encoder = nn.Sequential(
523+ nn.Conv2d(embed_dims, embed_dims, kernel_size=3, stride=1, padding=1, bias=False),
524+ nn.BatchNorm2d(embed_dims),
525+ nn.Conv2d(embed_dims, embed_dims, kernel_size=3, stride=2, padding=1, bias=False),
526+ nn.BatchNorm2d(embed_dims),
527+ nn.ReLU(),
528+ nn.AdaptiveAvgPool2d((1, 1)),
529+ )
530+ 
531+ def forward(self, feat):
532+ x = self.ego_feature_encoder(feat)
533+ x = x.flatten(1)
534+ return x
535+ 
536+ 
537+class SinusoidalPosEmb(nn.Module):
538+ def __init__(self, dim=256):
539+ super().__init__()
540+ self.dim = dim
541+ 
542+ def forward(self, x):
543+ device = x.device
544+ half = self.dim // 2
545+ emb = math.log(10000) / (half - 1)
546+ emb = torch.exp(torch.arange(half, device=device) * -emb)
547+ emb = x[:, None] * emb[None, :]
548+ emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
549+ return emb
550+ 
551+ 
552+class V1TrajPooler(nn.Module):
553+ def __init__(self, embed_dims=256):
554+ super().__init__()
555+ self.kps_generator = TrajSparsePoint3DKeyPointsGenerator()
556+ self.proj_drop = nn.Dropout(0.0)
557+ self.camera_encoder = nn.Sequential(
558+ nn.Linear(12, 256),
559+ nn.ReLU(inplace=True),
560+ nn.LayerNorm(256),
561+ nn.Linear(256, 256),
562+ nn.ReLU(inplace=True),
563+ nn.LayerNorm(256),
564+ )
565+ self.weights_fc = nn.Linear(256, 960)
566+ self.output_proj = nn.Linear(256, 256)
567+ 
568+ def forward(self, query, fpn_feats):
569+ pooled = F.adaptive_avg_pool2d(fpn_feats[0], (1, 1)).flatten(1).unsqueeze(1)
570+ out = query + pooled
571+ out = self.output_proj(out)
572+ return out
573+ 
574+ 
575+class V1ModulationLayer(nn.Module):
576+ def __init__(self, embed_dims=256):
577+ super().__init__()
578+ self.scale_shift_mlp = nn.Sequential(
579+ nn.Mish(),
580+ nn.Linear(embed_dims, embed_dims * 2),
581+ )
582+ 
583+ def forward(self, x, cond):
584+ scale_shift = self.scale_shift_mlp(cond)
585+ scale, shift = torch.chunk(scale_shift, 2, dim=-1)
586+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
587+ 
588+ 
589+class V11MotionPlanningRefinementModule(nn.Module):
590+ def __init__(self, embed_dims=256, motion_dim=24, status_dim=10):
591+ super().__init__()
592+ self.motion_cls_branch = nn.Sequential(
593+ nn.Linear(embed_dims, embed_dims),
594+ nn.ReLU(inplace=True),
595+ nn.LayerNorm(embed_dims),
596+ nn.Linear(embed_dims, embed_dims),
597+ nn.ReLU(inplace=True),
598+ nn.LayerNorm(embed_dims),
599+ nn.Linear(embed_dims, 1),
600+ )
601+ self.motion_reg_branch = nn.Sequential(
602+ nn.Linear(embed_dims, embed_dims),
603+ nn.ReLU(),
604+ nn.Linear(embed_dims, embed_dims),
605+ nn.ReLU(),
606+ nn.Linear(embed_dims, motion_dim),
607+ )
608+ self.plan_status_branch = nn.Sequential(
609+ nn.Linear(embed_dims, embed_dims),
610+ nn.ReLU(),
611+ nn.Linear(embed_dims, embed_dims),
612+ nn.ReLU(),
613+ nn.Linear(embed_dims, status_dim),
614+ )
615+ 
616+ def forward(self, x):
617+ pooled = x.mean(dim=1)
618+ motion_cls = self.motion_cls_branch(pooled)
619+ motion_reg = self.motion_reg_branch(pooled)
620+ status = self.plan_status_branch(pooled)
621+ return motion_cls, motion_reg, status
622+ 
623+ 
624+class V4DiffMotionPlanningRefinementModule(nn.Module):
625+ def __init__(self, embed_dims=256, plan_dim=12):
626+ super().__init__()
627+ self.plan_cls_branch = nn.Sequential(
628+ nn.Linear(embed_dims, embed_dims),
629+ nn.ReLU(inplace=True),
630+ nn.LayerNorm(embed_dims),
631+ nn.Linear(embed_dims, embed_dims),
632+ nn.ReLU(inplace=True),
633+ nn.LayerNorm(embed_dims),
634+ nn.Linear(embed_dims, 1),
635+ )
636+ self.plan_reg_branch = nn.Sequential(
637+ nn.Linear(embed_dims, embed_dims),
638+ nn.ReLU(),
639+ nn.Linear(embed_dims, embed_dims),
640+ nn.ReLU(),
641+ nn.Linear(embed_dims, plan_dim),
642+ )
643+ 
644+ def forward(self, x):
645+ pooled = x.mean(dim=1)
646+ plan_cls = self.plan_cls_branch(pooled)
647+ plan_reg = self.plan_reg_branch(pooled)
648+ return plan_cls, plan_reg
649+ 
650+ 
651+class V13MotionPlanningHeadLike(nn.Module):
652+ def __init__(self, embed_dims=256):
653+ super().__init__()
654+ self.instance_queue = InstanceQueue(embed_dims=embed_dims)
655+ 
656+ self.interact_layers = nn.ModuleList([
657+ MultiheadAttention(embed_dims=512, num_heads=8),
658+ MultiheadFlashAttention(embed_dims=512, num_heads=8),
659+ nn.LayerNorm(256),
660+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
661+ nn.LayerNorm(256),
662+ AsymmetricFFN(in_dims=256, embed_dims=256, hidden_dims=512, pre_norm_dim=256),
663+ nn.LayerNorm(256),
664+ 
665+ MultiheadAttention(embed_dims=512, num_heads=8),
666+ MultiheadFlashAttention(embed_dims=512, num_heads=8),
667+ nn.LayerNorm(256),
668+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
669+ nn.LayerNorm(256),
670+ AsymmetricFFN(in_dims=256, embed_dims=256, hidden_dims=512, pre_norm_dim=256),
671+ nn.LayerNorm(256),
672+ 
673+ MultiheadAttention(embed_dims=512, num_heads=8),
674+ MultiheadFlashAttention(embed_dims=512, num_heads=8),
675+ nn.LayerNorm(256),
676+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
677+ nn.LayerNorm(256),
678+ AsymmetricFFN(in_dims=256, embed_dims=256, hidden_dims=512, pre_norm_dim=256),
679+ nn.LayerNorm(256),
680+ 
681+ V11MotionPlanningRefinementModule(embed_dims=256, motion_dim=24, status_dim=10),
682+ ])
683+ 
684+ self.diff_layers = nn.ModuleList([
685+ V1TrajPooler(embed_dims=256),
686+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
687+ nn.LayerNorm(256),
688+ MultiheadFlashAttention(embed_dims=512, num_heads=8),
689+ nn.LayerNorm(256),
690+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
691+ nn.LayerNorm(256),
692+ AsymmetricFFN(in_dims=256, embed_dims=256, hidden_dims=512, pre_norm_dim=256),
693+ nn.LayerNorm(256),
694+ V1ModulationLayer(embed_dims=256),
695+ V4DiffMotionPlanningRefinementModule(embed_dims=256, plan_dim=12),
696+ 
697+ V1TrajPooler(embed_dims=256),
698+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
699+ nn.LayerNorm(256),
700+ MultiheadFlashAttention(embed_dims=512, num_heads=8),
701+ nn.LayerNorm(256),
702+ MultiheadFlashAttention(embed_dims=256, num_heads=8),
703+ nn.LayerNorm(256),
704+ AsymmetricFFN(in_dims=256, embed_dims=256, hidden_dims=512, pre_norm_dim=256),
705+ nn.LayerNorm(256),
706+ V1ModulationLayer(embed_dims=256),
707+ V4DiffMotionPlanningRefinementModule(embed_dims=256, plan_dim=12),
708+ ])
709+ 
710+ self.fc_before = nn.Linear(256, 512, bias=False)
711+ self.fc_after = nn.Linear(512, 256, bias=False)
712+ 
713+ self.motion_anchor_encoder = nn.Sequential(
714+ nn.Linear(256, 256),
715+ nn.ReLU(inplace=True),
716+ nn.LayerNorm(256),
717+ nn.Linear(256, 256),
718+ )
719+ self.plan_anchor_encoder = nn.Sequential(
720+ nn.Linear(256, 256),
721+ nn.ReLU(inplace=True),
722+ nn.LayerNorm(256),
723+ nn.Linear(256, 256),
724+ )
725+ self.plan_pos_encoder = nn.Sequential(
726+ nn.Linear(768, 256),
727+ nn.ReLU(inplace=True),
728+ nn.LayerNorm(256),
729+ nn.Linear(256, 256),
730+ )
731+ self.kps_generator = TrajSparsePoint3DKeyPointsGenerator()
732+ self.time_mlp = nn.Sequential(
733+ SinusoidalPosEmb(256),
734+ nn.Linear(256, 1024),
735+ nn.Mish(),
736+ nn.Linear(1024, 256),
737+ )
738+ 
739+ def forward(self, det_out, map_out, fpn_feats, metas=None):
740+ b = fpn_feats[0].shape[0]
741+ device = fpn_feats[0].device
742+ 
743+ det_query = det_out["query"]
744+ map_query = map_out["query"]
745+ 
746+ ego_feat = self.instance_queue(fpn_feats[0])
747+ ego_query = ego_feat.unsqueeze(1)
748+ 
749+ motion_query = ego_query + det_query.mean(dim=1, keepdim=True)
750+ motion_query = motion_query + self.motion_anchor_encoder(motion_query)
751+ 
752+ x = motion_query
753+ motion_cls = motion_reg = plan_status = None
754+ 
755+ for layer in self.interact_layers:
756+ if isinstance(layer, MultiheadAttention):
757+ x = self.fc_after(layer(self.fc_before(x)))
758+ elif isinstance(layer, MultiheadFlashAttention):
759+ if x.shape[-1] == 256 and layer.attn.embed_dim == 512:
760+ x = self.fc_after(layer(self.fc_before(x)))
761+ else:
762+ x = layer(x)
763+ elif isinstance(layer, nn.LayerNorm):
764+ x = layer(x)
765+ elif isinstance(layer, AsymmetricFFN):
766+ x = layer(x)
767+ elif isinstance(layer, V11MotionPlanningRefinementModule):
768+ motion_cls, motion_reg, plan_status = layer(x)
769+ 
770+ t = torch.arange(b, device=device).float()
771+ time_emb = self.time_mlp(t)
772+ 
773+ plan_seed = torch.cat([
774+ det_query.mean(dim=1),
775+ map_query.mean(dim=1),
776+ ego_feat
777+ ], dim=-1)
778+ plan_query = self.plan_pos_encoder(plan_seed).unsqueeze(1)
779+ plan_query = plan_query + self.plan_anchor_encoder(plan_query)
780+ 
781+ y = plan_query
782+ plan_cls = plan_reg = None
783+ for layer in self.diff_layers:
784+ if isinstance(layer, V1TrajPooler):
785+ y = layer(y, fpn_feats)
786+ elif isinstance(layer, MultiheadFlashAttention):
787+ if y.shape[-1] == 256 and layer.attn.embed_dim == 512:
788+ y = self.fc_after(layer(self.fc_before(y)))
789+ else:
790+ y = layer(y)
791+ elif isinstance(layer, nn.LayerNorm):
792+ y = layer(y)
793+ elif isinstance(layer, AsymmetricFFN):
794+ y = layer(y)
795+ elif isinstance(layer, V1ModulationLayer):
796+ y = layer(y, time_emb)
797+ elif isinstance(layer, V4DiffMotionPlanningRefinementModule):
798+ plan_cls, plan_reg = layer(y)
799+ 
800+ return {
801+ "motion_query": x,
802+ "plan_query": y,
803+ "motion_cls": motion_cls,
804+ "motion_reg": motion_reg,
805+ "plan_status": plan_status,
806+ "plan_cls": plan_cls,
807+ "plan_reg": plan_reg,
808+ }
809+ 
810+ 
811+class V1SparseDriveHead(nn.Module):
812+ def __init__(
813+ self,
814+ num_det_classes=10,
815+ num_map_classes=3,
816+ det_num_queries=300,
817+ map_num_queries=100,
818+ ):
819+ super().__init__()
820+ 
821+ self.det_head = Sparse4DHeadLike(
822+ mode="det",
823+ embed_dims=256,
824+ num_queries=det_num_queries,
825+ num_classes=num_det_classes,
826+ num_decoder=6,
827+ )
828+ self.map_head = Sparse4DHeadLike(
829+ mode="map",
830+ embed_dims=256,
831+ num_queries=map_num_queries,
832+ num_classes=num_map_classes,
833+ num_decoder=6,
834+ )
835+ self.motion_plan_head = V13MotionPlanningHeadLike(embed_dims=256)
836+ 
837+ def forward(self, fpn_feats, metas=None):
838+ det_out = self.det_head(fpn_feats, metas)
839+ map_out = self.map_head(fpn_feats, metas)
840+ motion_plan_out = self.motion_plan_head(det_out, map_out, fpn_feats, metas)
841+ 
842+ return {
843+ "det": det_out,
844+ "map": map_out,
845+ "motion_plan": motion_plan_out,
846+ }
847+ 
848+ 
849+class DenseDepthNetLike(nn.Module):
850+ def __init__(self, num_depth_layers=3, in_channels=256):
851+ super().__init__()
852+ self.depth_layers = nn.ModuleList([
853+ nn.Conv2d(in_channels, 1, kernel_size=1, stride=1, padding=0)
854+ for _ in range(num_depth_layers)
855+ ])
856+ 
857+ def forward(self, feats):
858+ outs = []
859+ for i, layer in enumerate(self.depth_layers):
860+ feat = feats[i]
861+ outs.append(layer(feat))
862+ return outs
863+ 
864+ 
865+class SimpleV1SparseDrive(BaseModel):
866+ def __init__(
867+ self,
868+ num_det_classes=10,
869+ num_map_classes=3,
870+ backbone_pretrained=None,
871+ use_grid_mask=True,
872+ data_preprocessor=None,
873+ init_cfg=None,
874+ ):
875+ super().__init__(data_preprocessor=data_preprocessor, init_cfg=init_cfg)
876+ 
877+ self.num_det_classes = num_det_classes
878+ self.num_map_classes = num_map_classes
879+ 
880+ self.img_backbone = ResNet(
881+ depth=50,
882+ num_stages=4,
883+ out_indices=(0, 1, 2, 3),
884+ frozen_stages=-1,
885+ norm_cfg=dict(type='BN', requires_grad=True),
886+ norm_eval=False,
887+ style='pytorch',
888+ init_cfg=(
889+ dict(type='Pretrained', checkpoint=backbone_pretrained)
890+ if backbone_pretrained is not None else None
891+ ),
892+ )
893+ 
894+ self.img_neck = FPN(
895+ in_channels=[256, 512, 1024, 2048],
896+ out_channels=256,
897+ num_outs=4,
898+ )
899+ 
900+ self.head = V1SparseDriveHead(
901+ num_det_classes=num_det_classes,
902+ num_map_classes=num_map_classes,
903+ det_num_queries=300,
904+ map_num_queries=100,
905+ )
906+ 
907+ self.depth_branch = DenseDepthNetLike(num_depth_layers=3, in_channels=256)
908+ self.grid_mask = GridMask() if use_grid_mask else nn.Identity()
909+ 
910+ def extract_feat(self, x):
911+ x = self.grid_mask(x)
912+ feats = self.img_backbone(x)
913+ feats = self.img_neck(feats)
914+ return feats
915+ 
916+ def _forward_impl(self, inputs, data_samples=None):
917+ feats = self.extract_feat(inputs)
918+ head_out = self.head(feats, data_samples)
919+ depth_out = self.depth_branch(feats)
920+ 
921+ outputs = {
922+ "head": head_out,
923+ "depth_branch": depth_out,
924+ "fpn_feats": feats,
925+ }
926+ return outputs
927+ 
928+ def _parse_data_samples(self, data_samples, device):
929+ parsed = {
930+ "gt_det_boxes": [],
931+ "gt_det_labels": [],
932+ "gt_map_points": [],
933+ "gt_map_labels": [],
934+ "gt_motion": [],
935+ "gt_plan": [],
936+ }
937+ 
938+ if data_samples is None:
939+ return parsed
940+ 
941+ for sample in data_samples:
942+ parsed["gt_det_boxes"].append(sample.get("gt_det_boxes", None))
943+ parsed["gt_det_labels"].append(sample.get("gt_det_labels", None))
944+ parsed["gt_map_points"].append(sample.get("gt_map_points", None))
945+ parsed["gt_map_labels"].append(sample.get("gt_map_labels", None))
946+ parsed["gt_motion"].append(sample.get("gt_motion", None))
947+ parsed["gt_plan"].append(sample.get("gt_plan", None))
948+ 
949+ return parsed
950+ 
951+ def loss(self, inputs, data_samples=None):
952+ outputs = self._forward_impl(inputs, data_samples)
953+ 
954+ det_out = outputs["head"]["det"]
955+ map_out = outputs["head"]["map"]
956+ motion_out = outputs["head"]["motion_plan"]
957+ depth_out = outputs["depth_branch"]
958+ 
959+ device = inputs.device
960+ b = inputs.shape[0]
961+ parsed_samples = self._parse_data_samples(data_samples, device)
962+ 
963+ loss_dict = {}
964+ 
965+ det_last_cls = det_out["all_cls_scores"][-1]
966+ det_last_reg = det_out["all_reg_preds"][-1]
967+ det_last_quality = det_out["all_quality_scores"][-1]
968+ 
969+ det_target_cls = torch.zeros(
970+ det_last_cls.shape[0], det_last_cls.shape[1],
971+ dtype=torch.long, device=device
972+ )
973+ det_target_reg = torch.zeros_like(det_last_reg)
974+ det_target_quality = torch.zeros_like(det_last_quality)
975+ 
976+ loss_dict["loss_det_cls"] = F.cross_entropy(
977+ det_last_cls.reshape(-1, det_last_cls.shape[-1]),
978+ det_target_cls.reshape(-1)
979+ )
980+ loss_dict["loss_det_reg"] = F.l1_loss(det_last_reg, det_target_reg)
981+ loss_dict["loss_det_quality"] = F.l1_loss(det_last_quality, det_target_quality)
982+ 
983+ map_last_cls = map_out["all_cls_scores"][-1]
984+ map_last_reg = map_out["all_reg_preds"][-1]
985+ 
986+ map_target_cls = torch.zeros(
987+ map_last_cls.shape[0], map_last_cls.shape[1],
988+ dtype=torch.long, device=device
989+ )
990+ map_target_reg = torch.zeros_like(map_last_reg)
991+ 
992+ loss_dict["loss_map_cls"] = F.cross_entropy(
993+ map_last_cls.reshape(-1, map_last_cls.shape[-1]),
994+ map_target_cls.reshape(-1)
995+ )
996+ loss_dict["loss_map_reg"] = F.l1_loss(map_last_reg, map_target_reg)
997+ 
998+ motion_cls = motion_out["motion_cls"]
999+ motion_reg = motion_out["motion_reg"]
1000+ plan_status = motion_out["plan_status"]
1001+ plan_cls = motion_out["plan_cls"]
1002+ plan_reg = motion_out["plan_reg"]
1003+ 
1004+ if data_samples is not None and len(parsed_samples["gt_motion"]) == b and parsed_samples["gt_motion"][0] is not None:
1005+ gt_motion = torch.stack([x.to(device) for x in parsed_samples["gt_motion"]], dim=0)
1006+ else:
1007+ gt_motion = torch.zeros_like(motion_reg)
1008+ 
1009+ if data_samples is not None and len(parsed_samples["gt_plan"]) == b and parsed_samples["gt_plan"][0] is not None:
1010+ gt_plan = torch.stack([x.to(device) for x in parsed_samples["gt_plan"]], dim=0)
1011+ else:
1012+ gt_plan = torch.zeros_like(plan_reg)
1013+ 
1014+ loss_dict["loss_motion_cls"] = F.binary_cross_entropy_with_logits(
1015+ motion_cls, torch.zeros_like(motion_cls)
1016+ )
1017+ loss_dict["loss_motion_reg"] = F.l1_loss(motion_reg, gt_motion)
1018+ loss_dict["loss_plan_status"] = F.cross_entropy(
1019+ plan_status, torch.zeros(b, dtype=torch.long, device=device)
1020+ )
1021+ loss_dict["loss_plan_cls"] = F.binary_cross_entropy_with_logits(
1022+ plan_cls, torch.zeros_like(plan_cls)
1023+ )
1024+ loss_dict["loss_plan_reg"] = F.l1_loss(plan_reg, gt_plan)
1025+ 
1026+ for i, d in enumerate(depth_out):
1027+ loss_dict[f"loss_depth_{i}"] = d.abs().mean()
1028+ 
1029+ return loss_dict
1030+ 
1031+ def predict(self, inputs, data_samples=None):
1032+ return self._forward_impl(inputs, data_samples)
1033+ 
1034+ def _forward(self, inputs, data_samples=None):
1035+ return self._forward_impl(inputs, data_samples)
1036+ 
1037+ def forward(self, inputs, data_samples=None, mode='tensor'):
1038+ if mode == 'loss':
1039+ return self.loss(inputs, data_samples)
1040+ elif mode == 'predict':
1041+ return self.predict(inputs, data_samples)
1042+ elif mode == 'tensor':
1043+ return self._forward(inputs, data_samples)
1044+ else:
1045+ raise ValueError(f'Invalid mode: {mode}')
Abenchmarks/v1sparsedrive/requirement.txt+6-0
@@ -0,0 +1,6 @@
1+mmdet==2.28.2
2+mmengine==0.10.7
3+scipy==1.13.1
4+opencv-python==4.7.0.68
5+pandas==2.3.3
6+numpy==1.23.4
Abenchmarks/v1sparsedrive/train.py+280-0
@@ -0,0 +1,280 @@
1+import os
2+import argparse
3+import random
4+import time
5+import numpy as np
6+import torch
7+from torch.utils.data import DataLoader
8+ 
9+import torch_npu
10+ 
11+from config import CONFIG
12+from datasets.dummy_dataset import DummySparseDriveDataset, simple_collate
13+from models import SimpleV1SparseDrive
14+ 
15+def patch_torch_npu_dvm_decomp_excludes(verbose: bool = True):
16+ """
17+ Runtime patch for:
18+ .../site-packages/torch_npu/_inductor/dvm/decomp.py
19+ 
20+ It appends:
21+ aten._native_batch_norm_legit_functional
22+ aten.native_batch_norm_backward
23+ 
24+ into decomps_to_exclude_npu (idempotent), then triggers patch_decomp().
25+ """
26+ import importlib
27+ import torch
28+ 
29+ mod = importlib.import_module("torch_npu._inductor.dvm.decomp")
30+ 
31+ aten = torch.ops.aten
32+ to_add = [
33+ aten._native_batch_norm_legit_functional,
34+ aten.native_batch_norm_backward,
35+ ]
36+ 
37+ if not hasattr(mod, "decomps_to_exclude_npu"):
38+ raise RuntimeError("torch_npu._inductor.dvm.decomp has no decomps_to_exclude_npu")
39+ 
40+ lst = mod.decomps_to_exclude_npu
41+ added = 0
42+ for op in to_add:
43+ if op not in lst:
44+ lst.append(op)
45+ added += 1
46+ 
47+ if verbose:
48+ print(
49+ f"[PATCH] torch_npu dvm decomp: added {added}/{len(to_add)} excludes. "
50+ f"Now len={len(lst)}"
51+ )
52+ if hasattr(mod, "patch_decomp"):
53+ mod.patch_decomp()
54+ if verbose:
55+ print("[PATCH] torch_npu dvm decomp: patch_decomp() called.")
56+ else:
57+ if verbose:
58+ print("[PATCH] torch_npu dvm decomp: no patch_decomp() found; only list updated.")
59+ 
60+ return mod
61+ 
62+def parse_args():
63+ parser = argparse.ArgumentParser()
64+ parser.add_argument(
65+ "--compile",
66+ action="store_true",
67+ help="Enable torch.compile (default: eager).",
68+ )
69+ parser.add_argument(
70+ "--profile",
71+ action="store_true",
72+ help="Enable torch_npu profiler (default: off).",
73+ )
74+ parser.add_argument(
75+ "--prof-dir",
76+ type=str,
77+ default=os.path.join(".", "prof", "example.prof"),
78+ help="Profiler output dir (only used when --profile).",
79+ )
80+ parser.add_argument(
81+ "--npu-backend",
82+ type=str,
83+ default="dvm")
84+ return parser.parse_args()
85+ 
86+ 
87+def set_seed(seed=3407):
88+ random.seed(seed)
89+ np.random.seed(seed)
90+ torch.manual_seed(seed)
91+ if torch.cuda.is_available():
92+ torch.cuda.manual_seed_all(seed)
93+ if hasattr(torch, "npu") and torch.npu.is_available():
94+ torch.npu.manual_seed_all(seed)
95+ 
96+ 
97+def get_profiler_config():
98+ return torch_npu.profiler._ExperimentalConfig(
99+ export_type=[
100+ torch_npu.profiler.ExportType.Text,
101+ torch_npu.profiler.ExportType.Db,
102+ ],
103+ profiler_level=torch_npu.profiler.ProfilerLevel.Level2,
104+ msprof_tx=False,
105+ aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
106+ l2_cache=False,
107+ op_attr=False,
108+ data_simplification=False,
109+ record_op_args=False,
110+ gc_detect_threshold=None,
111+ )
112+ 
113+ 
114+class NullProfiler:
115+ """兼容 profile 的空实现：start/step/stop 都是 no-op。"""
116+ def start(self): pass
117+ def step(self): pass
118+ def stop(self): pass
119+ 
120+ 
121+def maybe_enable_compile(model, enabled: bool):
122+ patch_torch_npu_dvm_decomp_excludes(verbose=True)
123+ if not enabled:
124+ print("[INFO] torch.compile disabled (eager mode).")
125+ return model
126+ try:
127+ model.img_backbone = torch.compile(
128+ model.img_backbone, dynamic=False, mode="reduce-overhead"
129+ )
130+ print("[INFO] torch.compile enabled successfully.")
131+ except Exception:
132+ print("[WARN] torch.compile failed, fallback to eager mode.")
133+ 
134+ return model
135+ 
136+ 
137+def build_profiler(enabled: bool, prof_output_dir: str, warmup_step_num=5, exec_step_num=10):
138+ if not enabled:
139+ print("[INFO] profiler disabled.")
140+ return NullProfiler()
141+ 
142+ os.makedirs(os.path.dirname(prof_output_dir), exist_ok=True)
143+ g_prof_config = get_profiler_config()
144+ 
145+ prof = torch_npu.profiler.profile(
146+ activities=[
147+ torch_npu.profiler.ProfilerActivity.CPU,
148+ torch_npu.profiler.ProfilerActivity.NPU,
149+ ],
150+ schedule=torch_npu.profiler.schedule(
151+ wait=0, warmup=warmup_step_num, active=exec_step_num, repeat=1
152+ ),
153+ on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(prof_output_dir),
154+ record_shapes=False,
155+ profile_memory=False,
156+ with_stack=False,
157+ with_modules=False,
158+ with_flops=False,
159+ experimental_config=g_prof_config,
160+ )
161+ print(f"[INFO] profiler enabled. output_dir = {prof_output_dir}")
162+ return prof
163+ 
164+ 
165+def main():
166+ args = parse_args()
167+ os.environ['TORCHINDUCTOR_NPU_BACKEND']=args.npu_backend
168+ torch.use_deterministic_algorithms(True)
169+ cfg = CONFIG
170+ set_seed(cfg["seed"])
171+ 
172+ device = (
173+ "npu" if hasattr(torch, "npu") and torch.npu.is_available()
174+ else "cuda" if torch.cuda.is_available()
175+ else "cpu"
176+ )
177+ print(f"[INFO] device = {device}")
178+ print(f"[INFO] flags: compile={args.compile}, profile={args.profile}")
179+ 
180+ dataset = DummySparseDriveDataset(**cfg["dummy_dataset"])
181+ dataloader = DataLoader(
182+ dataset,
183+ batch_size=cfg["batch_size"],
184+ shuffle=True,
185+ num_workers=cfg["num_workers"],
186+ collate_fn=simple_collate,
187+ drop_last=True,
188+ )
189+ 
190+ model = SimpleV1SparseDrive(**cfg["model"]).to(device)
191+ print(model)
192+ 
193+ # optional compile
194+ model = maybe_enable_compile(model, enabled=args.compile)
195+ 
196+ # optional profiler
197+ warmup_step_num = 5
198+ exec_step_num = 10
199+ prof = build_profiler(
200+ enabled=args.profile,
201+ prof_output_dir=args.prof_dir,
202+ warmup_step_num=warmup_step_num,
203+ exec_step_num=exec_step_num,
204+ )
205+ 
206+ model.train()
207+ optimizer = torch.optim.AdamW(
208+ model.parameters(),
209+ lr=cfg["lr"],
210+ weight_decay=cfg["weight_decay"],
211+ )
212+ 
213+ prof.start()
214+ print("[INFO] start training debug...")
215+ 
216+ for step, batch in enumerate(dataloader):
217+ if step >= cfg["train_steps"]:
218+ break
219+ 
220+ step_start_time = time.time()
221+ 
222+ inputs = batch["inputs"].to(device)
223+ data_samples = batch["data_samples"]
224+ 
225+ optimizer.zero_grad()
226+ loss_dict = model(inputs, data_samples=data_samples, mode="loss")
227+ total_loss = sum(v for v in loss_dict.values())
228+ total_loss.backward()
229+ optimizer.step()
230+ 
231+ prof.step()
232+ 
233+ step_time = time.time() - step_start_time
234+ log_str = f"[step {step}] total_loss={total_loss.item():.6f}, time={step_time:.4f}s"
235+ for k, v in loss_dict.items():
236+ log_str += f", {k}={v.item():.6f}"
237+ print(log_str)
238+ 
239+ print("[INFO] train debug finished.")
240+ prof.stop()
241+ 
242+ # quick sanity forward
243+ model.eval()
244+ with torch.no_grad():
245+ batch = next(iter(dataloader))
246+ inputs = batch["inputs"].to(device)
247+ data_samples = batch["data_samples"]
248+ outputs = model(inputs, data_samples=data_samples, mode="predict")
249+ 
250+ print("=" * 120)
251+ print("[det]")
252+ print("det query:", outputs["head"]["det"]["query"].shape)
253+ print("det cls layers:", len(outputs["head"]["det"]["all_cls_scores"]))
254+ print("det reg layers:", len(outputs["head"]["det"]["all_reg_preds"]))
255+ print("det last cls:", outputs["head"]["det"]["all_cls_scores"][-1].shape)
256+ print("det last reg:", outputs["head"]["det"]["all_reg_preds"][-1].shape)
257+ 
258+ print("\n[map]")
259+ print("map query:", outputs["head"]["map"]["query"].shape)
260+ print("map cls layers:", len(outputs["head"]["map"]["all_cls_scores"]))
261+ print("map reg layers:", len(outputs["head"]["map"]["all_reg_preds"]))
262+ print("map last cls:", outputs["head"]["map"]["all_cls_scores"][-1].shape)
263+ print("map last reg:", outputs["head"]["map"]["all_reg_preds"][-1].shape)
264+ 
265+ print("\n[motion_plan]")
266+ mp = outputs["head"]["motion_plan"]
267+ for k, v in mp.items():
268+ print(k, v.shape if v is not None else None)
269+ 
270+ print("\n[depth_branch]")
271+ for i, d in enumerate(outputs["depth_branch"]):
272+ print(f"depth[{i}] shape: {d.shape}")
273+ 
274+ print("\n[fpn_feats]")
275+ for i, f in enumerate(outputs["fpn_feats"]):
276+ print(f"fpn_feats[{i}] shape: {f.shape}")
277+ 
278+ 
279+if __name__ == "__main__":
280+ main()