diff --git a/ml/common/kin_torch.py b/ml/common/kin_torch.py index 91fb176..31587a1 100644 --- a/ml/common/kin_torch.py +++ b/ml/common/kin_torch.py @@ -118,12 +118,41 @@ class Kin: 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): """返回 dict: gm [B,49]:组对最小间隙 min(gap, cap)(gA*7+gB,gAbki', 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 B = q.shape[0] with torch.no_grad():