cbb/Robot_Arm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(kin/diag): kin.gaps 粗筛段也按点分块（COARSE_ENTRIES=3e7，≤446 点/块；实测上一版粗筛 4096 点仍 9.3 GiB）；diag_cl 加 D 段（子目标头精度 / 不用子目标 / 老师真实途经点）与目标误差进度比；rollout 加 use_sub、sub_override；本机执行清单更新

Browse Source
This commit is contained in:
Hark
2026-10-18 18:32:45 +00:00
parent 95ba43a1c7
commit 0653cf0d8d
4 changed files with 102 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/common/kin_torch.py
+36 -15
View File
@@ -118,11 +118,22 @@ class Kin:
return a, b
# ---------------- 间隙 ----------------
# 2026-10-19（用户实测 dagger/诊断吃满 24 GB 显存）：粗筛后幸存胶囊对数随 点数 × cap 暴涨，
# 一次几千个点、cap=40 mm 时中间张量可达十几 GB（训练里 eval 峰值 15.9 GiB 也是它）。
# 现在先按簇级粗筛估出每个点的幸存胶囊对数，按 PAIR_BUDGET 把 batch 切块逐块算，峰值与调用方的 batch 大小无关。
# 2026-10-19（用户实测 dagger/诊断吃满 24 GB 显存）：本模型 1936 条胶囊 / 491 组 / 67256 粗对 / 104.7 万细对。
# 显存有两段峰值，都随点数线性涨：
# 粗筛段 cc[:,cp_a]-cc[:,cp_b] 的广播中间量 ≈ 36 B × 点数 × 粗对数（4096 点实测 9.3 GiB，与 cap 无关）；
# 精算段 ≈ 每对几百字节 × 幸存细对数（4096 点 cap 40 幸存 2.4 亿对；旧版不分块实测 36.9 GiB）。
# 所以先按 COARSE_ENTRIES 把点切块（粗筛每块 ≤ ~1 GiB），块内再按幸存细对数 PAIR_BUDGET 切（精算每块 ≤ ~1 GiB）。
# 两段不同时存在，峰值与调用方 batch 大小无关。结果与不分块逐位一致（梯度经 cat 保留）。
COARSE_ENTRIES = 30_000_000
PAIR_BUDGET = 2_000_000
@staticmethod
def _cat_gaps(parts, starts):
return {'gm': torch.cat([p['gm'] for p in parts], 0),
'b': torch.cat([p['b'] + s for p, s in zip(parts, starts)]),
'i': torch.cat([p['i'] for p in parts]), 'j': torch.cat([p['j'] for p in parts]),
'gap': torch.cat([p['gap'] for p in parts])}
def gaps(self, q, cap=40.0, xf=None, max_pairs=None):
"""返回 dict：
gm [B,49]：组对最小间隙 min(gap, cap)（gA*7+gB，gA<gB；其余为 cap）
@@ -131,11 +142,25 @@ class Kin:
xf = xf if xf is not None else self.cum_xf(q)
Rs, Ts = xf
B = q.shape[0]
step = max(1, self.COARSE_ENTRIES // int(self.cp_a.shape[0]))
if B <= step:
return self._gaps_split(q, cap, xf)
st = list(range(0, B, step))
return self._cat_gaps([self._gaps_split(q[s:s + step], cap, (Rs[s:s + step], Ts[s:s + step])) for s in st], st)
def _coarse(self, q, cap, xf):
Rs, Ts = xf
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1) # [B,M]
del cc
return d2 <= (self.cp_lim + cap).pow(2)
def _gaps_split(self, q, cap, xf):
"""粗筛已在显存预算内：按幸存细对数再切块做精算"""
Rs, Ts = xf
B = q.shape[0]
with torch.no_grad():
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1)
per = ((d2 <= (self.cp_lim + cap).pow(2)).long() * self.cp_cnt).sum(1) # [B] 每点幸存胶囊对数
del cc, d2
per = (self._coarse(q, cap, xf).long() * self.cp_cnt).sum(1) # [B] 每点幸存细对数
if B <= 1 or int(per.sum()) <= self.PAIR_BUDGET:
bounds = [0, B]
else:
@@ -144,22 +169,18 @@ class Kin:
if c - base > self.PAIR_BUDGET and k > bounds[-1]:
bounds.append(k); base = cs[k - 1]
bounds.append(B)
del per
if len(bounds) == 2:
return self._gaps_one(q, cap, xf)
parts = [self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e])) for s, e in zip(bounds[:-1], bounds[1:])]
return {'gm': torch.cat([p['gm'] for p in parts], 0),
'b': torch.cat([p['b'] + s for p, s in zip(parts, bounds[:-1])]),
'i': torch.cat([p['i'] for p in parts]), 'j': torch.cat([p['j'] for p in parts]),
'gap': torch.cat([p['gap'] for p in parts])}
return self._cat_gaps([self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e])) for s, e in zip(bounds[:-1], bounds[1:])], bounds[:-1])
def _gaps_one(self, q, cap, xf):
Rs, Ts = xf
B = q.shape[0]
with torch.no_grad():
cc = torch.einsum('bkij,kj->bki', Rs[:, self.cl_g], self.cl_c.to(q.dtype)) + Ts[:, self.cl_g]
d2 = (cc[:, self.cp_a] - cc[:, self.cp_b]).pow(2).sum(-1) # [B,M]
keep = d2 <= (self.cp_lim + cap).pow(2)
keep = self._coarse(q, cap, xf)
bb, mm = keep.nonzero(as_tuple=True)
del keep
cnt = self.cp_cnt[mm]
rep_b = torch.repeat_interleave(bb, cnt)
off = torch.repeat_interleave(self.cp_start[mm] - torch.cumsum(cnt, 0) + cnt, cnt) + torch.arange(int(cnt.sum()), device=q.device)