Natural Memory NM2.1: 记忆路由器分叉、数据集缺陷修复与全轴评测证据
- 引入 MemoryRouterXL 与 v5/v6 流式多线程训练/编码管线 - 修复 prepare_memory_router_dataset 候选池重建缺陷(mega 家族 3568x 加速,输出逐字节相同) - 修复 v5 被破坏的拒答与多跳标签(train 未知样本 319 -> 16319,multi_hop 平均正例 1.00 -> 2.00) - 同存储预算下 V2-128 v6 逐轴 22/22 通过:Top-1 41.12% -> 94.62%,未知拒答 0.00% -> 100.00% - 记录三条被实测推翻的显然优化(logits_to_keep=1 反而慢 55%、XL 容量未带来收益) - 记忆手术跨架构可移植性 14/14,读写关闭时与原生模型逐位相同
This commit is contained in:
@@ -0,0 +1,49 @@
|
||||
"""Evaluate a saved dynamic-memory checkpoint."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
|
||||
from .model import DynamicMemoryConfig, DynamicMemoryLM
|
||||
from .tasks import sample_associative_batch
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint", default="V2_dpskw/checkpoints/latest.pt")
|
||||
parser.add_argument("--batches", type=int, default=100)
|
||||
parser.add_argument("--device", default="auto", choices=("auto", "cpu", "cuda"))
|
||||
args = parser.parse_args()
|
||||
|
||||
device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available() else ("cpu" if args.device == "auto" else args.device))
|
||||
checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
|
||||
config = DynamicMemoryConfig(**checkpoint["config"])
|
||||
model = DynamicMemoryLM(config).to(device)
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
model.eval()
|
||||
|
||||
for overwrite in (False, True):
|
||||
correct = 0
|
||||
total = 0
|
||||
with torch.no_grad():
|
||||
for _ in range(args.batches):
|
||||
batch = sample_associative_batch(
|
||||
batch_size=256,
|
||||
vocab_size=config.vocab_size,
|
||||
device=device,
|
||||
overwrite=overwrite,
|
||||
)
|
||||
memory = None
|
||||
for chunk in batch.learn_chunks:
|
||||
memory = model(chunk, memory=memory, update_memory=True).memory
|
||||
output = model(batch.query_input, memory=memory, update_memory=False)
|
||||
prediction = output.logits[:, 0].argmax(dim=-1)
|
||||
correct += int((prediction == batch.expected).sum())
|
||||
total += batch.expected.numel()
|
||||
print(f"overwrite={overwrite} accuracy={correct / total:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user