feat: 接口对齐 + FP16 量化(第一版优化方案)
- CTRUserDataset → CTRTestSeqDataset,构造参数对齐评测接口 - load_model 签名修正:ckpt_path 作为第一参数 - FP16 量化:model.half() + Embedding 保留 FP32 - move_batch_to_device 自动 FP32→FP16 转换 - 缓存时预转 FP16,减少推理循环开销 - requirements.txt 精简(去除 nvidia-* 包) - build_env.sh 标准化(set -e + pip install) - CLAUDE.md 更新开发命令、代码架构、关键接口说明
This commit is contained in:
+28
-13
@@ -118,15 +118,17 @@ def load_logids_from_file(file_path):
|
||||
return logids
|
||||
|
||||
|
||||
class CTRUserDataset(Dataset):
|
||||
"""按用户组织的 CTR 数据集"""
|
||||
class CTRTestSeqDataset(Dataset):
|
||||
"""按用户组织的 CTR 测试数据集(对齐评测接口)"""
|
||||
|
||||
def __init__(self, item_dict, user_seq=None, max_feasign_per_slot=None, pred_logids=None):
|
||||
def __init__(self, test_logids_ordered, item_dict, user_seq=None,
|
||||
max_feasign_per_slot=None, max_ctx_len=None):
|
||||
super().__init__()
|
||||
self.item_dict = item_dict
|
||||
self.user_seq = user_seq if user_seq else {}
|
||||
self.max_feasign_per_slot = max_feasign_per_slot
|
||||
self.pred_logids = pred_logids if pred_logids is not None else set()
|
||||
self.max_ctx_len = max_ctx_len
|
||||
self.pred_logids = set(test_logids_ordered) if test_logids_ordered else set()
|
||||
|
||||
self.user_items = defaultdict(list)
|
||||
for logid, rec in item_dict.items():
|
||||
@@ -236,7 +238,11 @@ def move_batch_to_device(batch, device):
|
||||
elif isinstance(batch, (list, tuple)):
|
||||
return [move_batch_to_device(x, device) for x in batch]
|
||||
elif torch.is_tensor(batch):
|
||||
return batch.to(device)
|
||||
x = batch.to(device)
|
||||
# 浮点 tensor → FP16,整数 tensor 保持不变
|
||||
if x.dtype == torch.float32:
|
||||
x = x.half()
|
||||
return x
|
||||
else:
|
||||
return batch
|
||||
|
||||
@@ -443,12 +449,12 @@ class CTRModel(nn.Module):
|
||||
# 模型加载入口
|
||||
# ============================================================
|
||||
|
||||
def load_model(device='cuda:0', ckpt_path=None):
|
||||
def load_model(ckpt_path, device='cuda:0'):
|
||||
"""加载模型并返回,供 evaluation.py 调用。
|
||||
|
||||
Args:
|
||||
ckpt_path: checkpoint 文件路径(评测系统传入 Path 对象)
|
||||
device: 推理设备(默认 'cuda:0')
|
||||
ckpt_path: checkpoint 文件路径,默认使用 infer.py 同目录下的 ckpt.pt
|
||||
|
||||
Returns:
|
||||
(model, device) 元组
|
||||
@@ -490,6 +496,11 @@ def load_model(device='cuda:0', ckpt_path=None):
|
||||
ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
|
||||
model.load_state_dict(ckpt['model_state_dict'])
|
||||
print(f"[INFO] Loaded checkpoint from {ckpt_path} (epoch={ckpt.get('epoch', '?')})")
|
||||
|
||||
# === FP16 量化:模型参数转半精度,Embedding 保留 FP32 ===
|
||||
model = model.half()
|
||||
model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32)
|
||||
print("[INFO] Model converted to FP16 (embedding kept in FP32)")
|
||||
else:
|
||||
print(f"[WARNING] Checkpoint {ckpt_path} not found, using random weights")
|
||||
|
||||
@@ -616,10 +627,11 @@ def main():
|
||||
print(f'[INFO] Test pred logids count: {len(test_pred_logids)}')
|
||||
|
||||
max_feasign_per_slot = {1: 2}
|
||||
test_dataset = CTRUserDataset(
|
||||
item_dict, user_seq,
|
||||
test_dataset = CTRTestSeqDataset(
|
||||
test_logids_ordered=list(test_pred_logids),
|
||||
item_dict=item_dict,
|
||||
user_seq=user_seq,
|
||||
max_feasign_per_slot=max_feasign_per_slot,
|
||||
pred_logids=test_pred_logids,
|
||||
)
|
||||
print(f'[INFO] num_users={test_dataset.num_users}, '
|
||||
f'total_samples={test_dataset.total_samples}, '
|
||||
@@ -634,9 +646,12 @@ def main():
|
||||
collate_fn=make_collate_fn(test_dataset.max_slot_id),
|
||||
)
|
||||
|
||||
# 收集 batches 并按分片缓存
|
||||
print('[INFO] collecting batches and saving sharded cache...')
|
||||
all_batches = [batch for batch in test_loader]
|
||||
# 收集 batches,预转 FP16 后按分片缓存
|
||||
print('[INFO] collecting batches (pre-converting to FP16) and saving sharded cache...')
|
||||
all_batches = []
|
||||
for batch in test_loader:
|
||||
batch = move_batch_to_device(batch, torch.device('cpu'))
|
||||
all_batches.append(batch)
|
||||
|
||||
batches_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
shard_idx = 0
|
||||
|
||||
Reference in New Issue
Block a user