fix(kin/diag): gaps 默认不拼胶囊对(pairs=True 仅 APG 用;实测 4096×40 峰值 4.46 GiB 全在拼接 8.6e7 对);diag_cl 删 D 段(子目标已排除)、加 E 段(开环幅度与抹 dq 敏感度、老师中段带速度/静止起步闭环);rollout 加 qp0;本机执行清单更新
This commit is contained in:
+16
-12
@@ -129,24 +129,27 @@ class Kin:
|
||||
|
||||
@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])}
|
||||
out = {'gm': torch.cat([p.pop('gm') for p in parts], 0)}
|
||||
if 'b' in parts[0]:
|
||||
for p, s in zip(parts, starts): p['b'] += s
|
||||
for k in ('b', 'i', 'j', 'gap'): out[k] = torch.cat([p.pop(k) for p in parts]) # 逐键拼、拼完即释放
|
||||
return out
|
||||
|
||||
def gaps(self, q, cap=40.0, xf=None, max_pairs=None):
|
||||
def gaps(self, q, cap=40.0, xf=None, max_pairs=None, pairs=False):
|
||||
"""返回 dict:
|
||||
gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gA<gB;其余为 cap)
|
||||
b, i, j, gap:幸存胶囊对(gap 可微)
|
||||
pairs=True 时另有 b, i, j, gap:幸存胶囊对(gap 可微;只有 APG 的 hinge 与 check_parity 需要)
|
||||
2026-10-19(用户实测 4096 点 cap 40 峰值 4.46 GiB 全在拼接 8.6e7 对的 b/i/j/gap):
|
||||
默认只返回 gm,每块算完就丢掉胶囊对,峰值只剩单块。
|
||||
"""
|
||||
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)
|
||||
return self._gaps_split(q, cap, xf, pairs)
|
||||
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)
|
||||
return self._cat_gaps([self._gaps_split(q[s:s + step], cap, (Rs[s:s + step], Ts[s:s + step]), pairs) for s in st], st)
|
||||
|
||||
def _coarse(self, q, cap, xf):
|
||||
Rs, Ts = xf
|
||||
@@ -155,7 +158,7 @@ class Kin:
|
||||
del cc
|
||||
return d2 <= (self.cp_lim + cap).pow(2)
|
||||
|
||||
def _gaps_split(self, q, cap, xf):
|
||||
def _gaps_split(self, q, cap, xf, pairs):
|
||||
"""粗筛已在显存预算内:按幸存细对数再切块做精算"""
|
||||
Rs, Ts = xf
|
||||
B = q.shape[0]
|
||||
@@ -171,10 +174,10 @@ class Kin:
|
||||
bounds.append(B)
|
||||
del per
|
||||
if len(bounds) == 2:
|
||||
return self._gaps_one(q, cap, xf)
|
||||
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])
|
||||
return self._gaps_one(q, cap, xf, pairs)
|
||||
return self._cat_gaps([self._gaps_one(q[s:e], cap, (Rs[s:e], Ts[s:e]), pairs) for s, e in zip(bounds[:-1], bounds[1:])], bounds[:-1])
|
||||
|
||||
def _gaps_one(self, q, cap, xf):
|
||||
def _gaps_one(self, q, cap, xf, pairs):
|
||||
Rs, Ts = xf
|
||||
B = q.shape[0]
|
||||
with torch.no_grad():
|
||||
@@ -195,6 +198,7 @@ class Kin:
|
||||
gap = d - self.pair_need[pidx].to(q.dtype)
|
||||
gm = torch.full((B * 49,), float(cap), dtype=q.dtype, device=q.device)
|
||||
gm = gm.scatter_reduce(0, rep_b * 49 + self.pair_gp[pidx], gap.clamp(max=cap), reduce='amin', include_self=True)
|
||||
if not pairs: return {'gm': gm.view(B, 49)}
|
||||
return {'gm': gm.view(B, 49), 'b': rep_b, 'i': pi, 'j': pj, 'gap': gap}
|
||||
|
||||
def min_gap(self, q, cap=40.0):
|
||||
|
||||
Reference in New Issue
Block a user