OwnerSunshine530
c0c23ad224
不同用户被因果mask隔离，过滤非测试用户对测试样本AUC/PCOC零影响。 流式加载只持有测试用户记录，避免 CTRTestSeqDataset 构造期 OOM； 过滤结果缓存到 bench_filtered_cache.pt，后续秒级复用。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
239 lines
9.3 KiB
Python
239 lines
9.3 KiB
Python
"""本地测量闭环：设置 infer.CONFIG，跑推理，同步计时，打印 AUC/PCOC/延迟/总分。
|
||
|
||
不进提交包。**以子进程方式运行**（AI Studio 内核禁止 import torch）：
|
||
|
||
%cd /home/aistudio/code
|
||
!python bench.py --smoke 50 # 冒烟：只跑前 50 batch
|
||
!python bench.py # 默认基线
|
||
!python bench.py --fp32 # FP32 天花板（Task 3）
|
||
!python bench.py --rebuild # 强制重建过滤缓存
|
||
|
||
关键设计——只保留“测试用户”的数据：
|
||
不同用户被因果 mask 完全隔离，非测试用户的前向输出不参与打分；过滤掉它们
|
||
对测试样本的 AUC/PCOC 没有任何影响，却能把数据量从 924 万条降到一小部分，
|
||
避免 CTRTestSeqDataset 构造时 OOM。过滤后的数据缓存到磁盘，后续秒级复用。
|
||
"""
|
||
import os
|
||
import sys
|
||
import time
|
||
from collections import defaultdict
|
||
from pathlib import Path
|
||
|
||
# baseline 把依赖装在 --target 目录（非默认 site-packages），import 前先加 sys.path
|
||
for _p in ("/home/aistudio/external-libraries", "/home/aistudio/libraries",
|
||
os.path.abspath("../libraries"), os.path.abspath("./libraries")):
|
||
if os.path.isdir(_p) and _p not in sys.path:
|
||
sys.path.insert(0, _p)
|
||
|
||
import numpy as np
|
||
import torch
|
||
from torch.utils.data import DataLoader
|
||
|
||
import infer # 同目录
|
||
|
||
|
||
def _test_user_ids(test_csv):
|
||
"""从 test.csv 读出所有测试用户 id（第 2 列 userid）。"""
|
||
users = set()
|
||
with open(test_csv) as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
parts = line.split(",")
|
||
if len(parts) >= 2:
|
||
users.add(int(parts[1]))
|
||
return users
|
||
|
||
|
||
def _load_filtered(history_dir, test_csv, test_users):
|
||
"""流式读取所有文件，只保留 userid ∈ test_users 的记录（不持有完整字典，防 OOM）。
|
||
|
||
解析逻辑与 infer.load_sample_files 完全一致，只是多了一道用户过滤。
|
||
"""
|
||
files = (sorted(history_dir.glob("*.csv")) if history_dir.exists() else []) + [test_csv]
|
||
print(f"[BENCH] 流式过滤加载 {len(files)} 个文件（仅保留 {len(test_users)} 个测试用户）...")
|
||
item_dict = {}
|
||
user_logs = defaultdict(list)
|
||
for fp in files:
|
||
has_clk = infer._detect_has_clk(fp)
|
||
min_parts = 5 if has_clk else 4
|
||
kept = 0
|
||
with open(fp) as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
parts = line.split(",")
|
||
if len(parts) < min_parts:
|
||
continue
|
||
userid = int(parts[1])
|
||
if userid not in test_users:
|
||
continue
|
||
logid = int(parts[0])
|
||
adid = int(parts[2])
|
||
if has_clk:
|
||
clk = int(parts[3])
|
||
timestamp = int(parts[4])
|
||
fs = 5
|
||
else:
|
||
clk = 0
|
||
timestamp = int(parts[3])
|
||
fs = 4
|
||
signs, slots = [], []
|
||
for pair in parts[fs:]:
|
||
if ":" in pair:
|
||
s, sl = pair.split(":", 1)
|
||
signs.append(int(s))
|
||
slots.append(int(sl))
|
||
item_dict[logid] = {
|
||
"logid": logid, "userid": userid, "adid": adid,
|
||
"clk": clk, "timestamp": timestamp,
|
||
"signs": np.array(signs, dtype=np.int64),
|
||
"slots": np.array(slots, dtype=np.int64),
|
||
}
|
||
user_logs[userid].append((timestamp, logid))
|
||
kept += 1
|
||
print(f" {fp.name}: has_clk={has_clk}, kept={kept}")
|
||
|
||
user_seq = {}
|
||
for u, logs in user_logs.items():
|
||
logs.sort(key=lambda x: x[0])
|
||
user_seq[u] = [lid for _, lid in logs]
|
||
print(f"[BENCH] 过滤后：{len(item_dict)} 条记录，{len(user_seq)} 个用户")
|
||
return item_dict, user_seq
|
||
|
||
|
||
def _get_data(cur, ref, rebuild=False):
|
||
"""取过滤后的 (item_dict, user_seq)，优先读磁盘缓存。"""
|
||
cache = cur / "bench_filtered_cache.pt"
|
||
test_csv = ref / "test.csv"
|
||
history = ref / "history"
|
||
if cache.exists() and not rebuild:
|
||
print(f"[BENCH] 读取过滤缓存：{cache}")
|
||
d = torch.load(cache, weights_only=False)
|
||
return d["item_dict"], d["user_seq"]
|
||
test_users = _test_user_ids(test_csv)
|
||
item_dict, user_seq = _load_filtered(history, test_csv, test_users)
|
||
torch.save({"item_dict": item_dict, "user_seq": user_seq}, cache)
|
||
print(f"[BENCH] 已缓存 -> {cache}")
|
||
return item_dict, user_seq
|
||
|
||
|
||
def run_once(config_override=None, batch_size=50, max_batches=None,
|
||
max_feasign_per_slot=None, rebuild=False):
|
||
"""跑一次本地推理并打分。返回 infer._cal_score 的结果 dict。"""
|
||
if config_override is None:
|
||
config_override = {}
|
||
if max_feasign_per_slot is None:
|
||
max_feasign_per_slot = {1: 2}
|
||
|
||
infer.CONFIG.update(config_override)
|
||
infer.CONFIG["sync_timing"] = True
|
||
|
||
cur = Path(__file__).parent
|
||
ref = cur / "dataset"
|
||
test_csv = ref / "test.csv"
|
||
label_file = ref / "label_data.txt"
|
||
|
||
# ----- 取数据（过滤+缓存）-----
|
||
item_dict, user_seq = _get_data(cur, ref, rebuild=rebuild)
|
||
test_logids = infer.load_logids_from_file(test_csv)
|
||
ds = infer.CTRTestSeqDataset(
|
||
test_logids_ordered=list(test_logids), item_dict=item_dict,
|
||
user_seq=user_seq, max_feasign_per_slot=max_feasign_per_slot, max_ctx_len=None,
|
||
)
|
||
loader = DataLoader(
|
||
ds, batch_size=batch_size, shuffle=False, num_workers=0,
|
||
collate_fn=infer.make_collate_fn(ds.max_slot_id),
|
||
)
|
||
batches = []
|
||
for b in loader:
|
||
batches.append(infer.move_batch_to_device(b, torch.device("cpu")))
|
||
if max_batches is not None and len(batches) >= max_batches:
|
||
break
|
||
|
||
# 释放构造期内存，降低推理峰值
|
||
del item_dict, user_seq, ds, loader
|
||
import gc
|
||
gc.collect()
|
||
|
||
# ----- 加载模型 -----
|
||
model, dev = infer.load_model(ckpt_path=None)
|
||
|
||
# ----- 推理 + 同步计时 -----
|
||
logid2p = {}
|
||
t_sum = 0.0
|
||
cuda = (dev.type == "cuda")
|
||
with torch.inference_mode():
|
||
for b in batches:
|
||
b = infer.move_batch_to_device(b, dev)
|
||
pm = b["pred_mask"].bool()
|
||
if cuda:
|
||
torch.cuda.synchronize()
|
||
t0 = time.time()
|
||
logits, _ = model(b)
|
||
probs = torch.sigmoid(logits.squeeze(-1))
|
||
if cuda:
|
||
torch.cuda.synchronize()
|
||
t_sum += time.time() - t0
|
||
for lid, p in zip(b["logid"][pm].cpu().tolist(), probs[pm].cpu().tolist()):
|
||
logid2p[lid] = p
|
||
|
||
# ----- 按 test.csv 顺序写 predict.txt 并打分 -----
|
||
order = [int(l.split(",")[0]) for l in open(test_csv) if l.strip()]
|
||
missing = [lid for lid in order if lid not in logid2p]
|
||
if missing:
|
||
print(f"[BENCH][WARN] {len(missing)} 个测试 logid 没预测到（前几个 {missing[:5]}）")
|
||
pred_path = cur / "predict.txt"
|
||
with open(pred_path, "w") as f:
|
||
for lid in order:
|
||
f.write(f"{logid2p.get(lid, 0.0)}\n")
|
||
|
||
res = infer._cal_score(pred_path, label_file, default_latency=t_sum)
|
||
print(
|
||
f"[BENCH] cfg={config_override} bs={batch_size}"
|
||
f"{'' if max_batches is None else f' (first {max_batches} batches)'}"
|
||
f" -> AUC={res['auc']:.5f} PCOC={res['pcoc']:.4f}"
|
||
f" lat={res['latency']:.2f}s score={res['score_all']:.2f}"
|
||
)
|
||
return res
|
||
|
||
|
||
def _parse_args():
|
||
import argparse
|
||
ap = argparse.ArgumentParser(description="CTI 推理测量闭环（子进程跑：!python bench.py ...）")
|
||
ap.add_argument("--smoke", type=int, default=None, help="只跑前 N 个 batch（冒烟）")
|
||
ap.add_argument("--bs", type=int, default=50, help="batch_size（本地参考）")
|
||
ap.add_argument("--fp32", action="store_true", help="FP32 天花板 = 关 fp16 + 关 expert 合并")
|
||
ap.add_argument("--no-fp16", action="store_true", help="关闭半精度")
|
||
ap.add_argument("--no-merge", action="store_true", help="关闭 expert 合并")
|
||
ap.add_argument("--signid", choices=["clamp", "modulo"], default=None, help="sign-id 处理方式")
|
||
ap.add_argument("--merge-th", type=float, default=None, help="expert 合并余弦阈值")
|
||
ap.add_argument("--keep", type=str, default=None,
|
||
help="逗号分隔的 keep_fp32_modules，如 linear,rep_encoder.input_norm")
|
||
ap.add_argument("--feasign-none", action="store_true",
|
||
help="不截断特征（max_feasign_per_slot=None）")
|
||
ap.add_argument("--rebuild", action="store_true", help="强制重建过滤缓存")
|
||
return ap.parse_args()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
a = _parse_args()
|
||
cfg = {}
|
||
if a.fp32:
|
||
cfg["fp16"] = False
|
||
cfg["expert_merge"] = False
|
||
if a.no_fp16:
|
||
cfg["fp16"] = False
|
||
if a.no_merge:
|
||
cfg["expert_merge"] = False
|
||
if a.signid:
|
||
cfg["signid_mode"] = a.signid
|
||
if a.merge_th is not None:
|
||
cfg["merge_threshold"] = a.merge_th
|
||
if a.keep is not None:
|
||
cfg["keep_fp32_modules"] = tuple(x for x in a.keep.split(",") if x)
|
||
mf = None if a.feasign_none else {1: 2}
|
||
run_once(cfg, batch_size=a.bs, max_batches=a.smoke, max_feasign_per_slot=mf, rebuild=a.rebuild)
|