fix(kin): gaps 按幸存胶囊对数(PAIR_BUDGET=2e6)自动分块,显存峰值与调用方 batch 无关
This commit is contained in:
@@ -118,12 +118,41 @@ class Kin:
|
|||||||
return a, b
|
return a, b
|
||||||
|
|
||||||
# ---------------- 间隙 ----------------
|
# ---------------- 间隙 ----------------
|
||||||
|
# 2026-10-19(用户实测 dagger/诊断吃满 24 GB 显存):粗筛后幸存胶囊对数随 点数 × cap 暴涨,
|
||||||
|
# 一次几千个点、cap=40 mm 时中间张量可达十几 GB(训练里 eval 峰值 15.9 GiB 也是它)。
|
||||||
|
# 现在先按簇级粗筛估出每个点的幸存胶囊对数,按 PAIR_BUDGET 把 batch 切块逐块算,峰值与调用方的 batch 大小无关。
|
||||||
|
PAIR_BUDGET = 2_000_000
|
||||||
|
|
||||||
def gaps(self, q, cap=40.0, xf=None, max_pairs=None):
|
def gaps(self, q, cap=40.0, xf=None, max_pairs=None):
|
||||||
"""返回 dict:
|
"""返回 dict:
|
||||||
gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gA<gB;其余为 cap)
|
gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gA<gB;其余为 cap)
|
||||||
b, i, j, gap:幸存胶囊对(gap 可微)
|
b, i, j, gap:幸存胶囊对(gap 可微)
|
||||||
"""
|
"""
|
||||||
xf = xf if xf is not None else self.cum_xf(q)
|
xf = xf if xf is not None else self.cum_xf(q)
|
||||||
|
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
|
||||||
|
if B <= 1 or int(per.sum()) <= self.PAIR_BUDGET:
|
||||||
|
bounds = [0, B]
|
||||||
|
else:
|
||||||
|
cs = torch.cumsum(per, 0).tolist(); bounds = [0]; base = 0
|
||||||
|
for k, c in enumerate(cs):
|
||||||
|
if c - base > self.PAIR_BUDGET and k > bounds[-1]:
|
||||||
|
bounds.append(k); base = cs[k - 1]
|
||||||
|
bounds.append(B)
|
||||||
|
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])}
|
||||||
|
|
||||||
|
def _gaps_one(self, q, cap, xf):
|
||||||
Rs, Ts = xf
|
Rs, Ts = xf
|
||||||
B = q.shape[0]
|
B = q.shape[0]
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
|
|||||||
Reference in New Issue
Block a user