第五轮:锁/续跑/超时/自动分支/R4 探针与 encode-once/控制台加固;阶段内断点续跑(gen 按条、prep 按 shard、训练与 APG 定时原子存档+体面停止);分歧以实现方意见为准
This commit is contained in:
@@ -126,6 +126,22 @@ class Policy(nn.Module):
|
||||
x = z - self.decode(mem, z, zero, one)
|
||||
return x, v, s
|
||||
|
||||
def sample_shared(self, feat1, z):
|
||||
"""部署接口(ONNX,2026-10-18 第五轮,R4 实测后改):一次观测、K 条候选。
|
||||
feat1 [1,F],z [K,16,6] → act [K,16,6],value [K],sub [K,13](value/sub 对 K 行相同,展开只为保持运行时读法)。
|
||||
此前部署把同一观测复制 K=8 行喂 sample(),encoder 被重复算 8 次(约占总计算 29%);这里只编码一次再展开给解码器。"""
|
||||
mem = self.encode(feat1)
|
||||
v, s = self.heads(mem)
|
||||
K = z.shape[0]
|
||||
if self.head == 'wta':
|
||||
act, _ = self.wta_modes(mem)
|
||||
idx = torch.arange(K, device=z.device) % self.K
|
||||
return act[0, idx] + 0.0 * z, v.expand(K), s.expand(K, -1)
|
||||
m = mem.expand(K, -1, -1)
|
||||
one = torch.ones(K, device=z.device, dtype=z.dtype)
|
||||
x = z - self.decode(m, z, torch.zeros_like(one), one)
|
||||
return x, v.expand(K), s.expand(K, -1)
|
||||
|
||||
def sample_k(self, feat, K, z=None):
|
||||
"""仿真 / APG 用:每个观测 K 条候选 → act [B,K,16,6], value [B], sub [B,13]"""
|
||||
mem = self.encode(feat)
|
||||
|
||||
Reference in New Issue
Block a user