[feat] model_config.summaries_set：训练期 PCOC 与 scalars TensorBoard 摘要

docs/source/models/evaluation_metrics.md

@@ -60,6 +60,45 @@ model_config {
   - decay_rate: 默认为0.9，历史指标的衰减率。
   - decay_step: 默认为100，历史指标每训练多少step会进行一次衰减，配置需要可以被train_config.log_step_count_steps整除。
 
+### 训练期 TensorBoard 摘要（summaries_set）
+
+在 `model_config` 中配置 `summaries_set`，可在训练过程中向 TensorBoard 写入 **PCOC** 与 **特征切片 scalars**（与 `train_metrics` 不同，后者用于在线 AUC 等指标）。摘要按 `train_config.log_step_count_steps` 间隔累积并在日志步写入。
+
+```
+model_config {
+    # PCOC -> summary/pcoc, summary/predicted_ctr, summary/observed_ctr
+    summaries_set {
+        pcoc {
+            pred_name: "probs"      # prediction_dict 中的 key；logits 会先 sigmoid
+            label_name: "clk"       # 可选，默认取 data_config 第一个 label
+            epsilon: 1e-7           # observed ctr 接近 0 时的除数稳定项
+        }
+    }
+    # 特征切片 -> summary/<name>
+    summaries_set {
+        scalars {
+            name: "id_1_bucket_pred"
+            pred_name: "probs"
+            feature_name: "id_1"
+            feature_value: "1005"
+        }
+    }
+}
+```
+
+- **pcoc**
+  - `pred_name`: 预测张量名，默认 `probs`
+  - `label_name`: 标签字段名；多任务时需显式指定
+  - `epsilon`: 默认 `1e-7`
+- **scalars**
+  - `name`: TensorBoard tag 后缀，写入 `summary/<name>`
+  - `pred_name`: 默认 `probs`
+  - `feature_name` / `feature_value`: 可选；同时配置时仅统计该特征取值的样本；省略则统计全 batch 均值
+
+> 分布式训练时，各 rank 在本地累积摘要，仅 **rank 0** 写入 TensorBoard，数值反映 rank 0 所见分片数据的统计，而非全局 all-reduce 结果（与当前 `train_metrics` 行为一致）。
+
+需开启 `train_config.use_tensorboard: true`（默认开启）。
+
 ______________________________________________________________________
 
 ## 指标详情

tzrec/main.py

@@ -240,6 +240,7 @@ def _log_train(
     plogger: Optional[ProgressLogger] = None,
     summary_writer: Optional[SummaryWriter] = None,
     train_metrics: Optional[Dict[str, torch.Tensor]] = None,
+    train_summaries: Optional[Dict[str, torch.Tensor]] = None,
 ) -> None:
     """Logging current training step."""
     if plogger is not None:
@@ -317,6 +318,9 @@ def _log_train(
         if train_metrics:
             for k, v in train_metrics.items():
                 summary_writer.add_scalar(f"metric/{k}", v, step)
+        if train_summaries:
+            for k, v in train_summaries.items():
+                summary_writer.add_scalar(k, v, step)
 
 
 def _train_and_evaluate(
@@ -491,8 +495,10 @@ def run_eval(step: int, epoch: int) -> None:
                     dataloader_state, batch.checkpoint_info
                 )
                 _model.update_train_metric(predictions, batch)
+                _model.update_train_summary(predictions, batch)
                 if i_step % train_config.log_step_count_steps == 0:
                     train_metrics = _model.compute_train_metric()
+                    train_summaries = _model.compute_train_summary()
                     _log_train(
                         i_step,
                         losses,
@@ -502,6 +508,7 @@ def run_eval(step: int, epoch: int) -> None:
                         plogger=plogger,
                         summary_writer=summary_writer,
                         train_metrics=train_metrics,
+                        train_summaries=train_summaries,
                     )
 
                 for lr in lr_scheduler:

tzrec/metrics/train_summary.py

@@ -0,0 +1,165 @@
+# Copyright (c) 2024, Alibaba Group;
+# Licensed under the Apache License, Version 2.0 (the "License");
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#    http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch import nn
+
+from tzrec.protos.summary_pb2 import ModelSummaries
+
+BASE_DATA_GROUP = "__BASE__"
+
+
+def resolve_pred_tensor(
+    predictions: Dict[str, torch.Tensor], pred_name: str
+) -> torch.Tensor:
+    if pred_name not in predictions:
+        raise ValueError(
+            f'summaries pred_name "{pred_name}" not found in prediction_dict keys: '
+            f"{sorted(predictions.keys())}"
+        )
+    pred = predictions[pred_name].float()
+    if pred_name == "logits" or pred_name.startswith("logits"):
+        pred = torch.sigmoid(pred)
+    if pred.dim() > 1:
+        if pred.dim() == 2 and pred.size(-1) == 2:
+            pred = pred[:, 1]
+        else:
+            pred = pred.reshape(-1)
+    return pred.reshape(-1)
+
+
+def _feature_values(batch: Any, feature_name: str) -> torch.Tensor:
+    if BASE_DATA_GROUP in batch.sparse_features:
+        sparse = batch.sparse_features[BASE_DATA_GROUP].to_dict()
+        if feature_name in sparse:
+            return sparse[feature_name].to_padded_dense(1)[:, 0]
+    if BASE_DATA_GROUP in batch.dense_features:
+        dense = batch.dense_features[BASE_DATA_GROUP].to_dict()
+        if feature_name in dense:
+            return dense[feature_name].values().reshape(-1)
+    raise ValueError(
+        f'summaries feature_name "{feature_name}" not found in batch features'
+    )
+
+
+def build_feature_mask(
+    batch: Any, feature_name: str, feature_value: str
+) -> torch.Tensor:
+    feat = _feature_values(batch, feature_name)
+    try:
+        target = float(feature_value)
+        if feat.is_floating_point() or feat.dtype in (torch.int32, torch.int64):
+            return feat.float() == target
+    except (TypeError, ValueError):
+        pass
+    # string / id features: compare as int when possible
+    try:
+        return feat.long() == int(float(feature_value))
+    except (TypeError, ValueError):
+        return feat == feat.new_tensor(feature_value)
+
+
+class _RunningMeanState(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.register_buffer("_sum", torch.zeros(1), persistent=False)
+        self.register_buffer("_count", torch.zeros(1), persistent=False)
+
+    def update(self, values: torch.Tensor, mask: Optional[torch.Tensor] = None) -> None:
+        values = values.reshape(-1).float()
+        if mask is not None:
+            values = values[mask.reshape(-1)]
+        if values.numel() == 0:
+            return
+        self._sum += values.sum()
+        self._count += values.numel()
+
+    def compute_mean(self) -> Optional[torch.Tensor]:
+        if self._count.item() == 0:
+            return None
+        return self._sum / self._count
+
+    def reset(self) -> None:
+        self._sum.zero_()
+        self._count.zero_()
+
+
+class TrainSummaryModule(nn.Module):
+    """Accumulate training summaries between log steps."""
+
+    def __init__(self, summary_cfg: ModelSummaries) -> None:
+        super().__init__()
+        self._summary_cfg = summary_cfg
+        self._state = _RunningMeanState()
+        self._label_state = _RunningMeanState()
+
+    @property
+    def summary_cfg(self) -> ModelSummaries:
+        """Configured summary proto for this module."""
+        return self._summary_cfg
+
+    def update(
+        self,
+        predictions: Dict[str, torch.Tensor],
+        batch: Any,
+        label_name: str,
+    ) -> None:
+        summary_type = self._summary_cfg.WhichOneof("summary")
+        if summary_type == "pcoc":
+            cfg = self._summary_cfg.pcoc
+            pred_name = cfg.pred_name or "probs"
+            preds = resolve_pred_tensor(predictions, pred_name)
+            label = batch.labels[label_name].float().reshape(-1)
+            self._state.update(preds)
+            self._label_state.update(label)
+        elif summary_type == "scalars":
+            cfg = self._summary_cfg.scalars
+            pred_name = cfg.pred_name or "probs"
+            preds = resolve_pred_tensor(predictions, pred_name)
+            if cfg.HasField("feature_name") and cfg.HasField("feature_value"):
+                mask = build_feature_mask(batch, cfg.feature_name, cfg.feature_value)
+                self._state.update(preds, mask)
+            else:
+                self._state.update(preds)
+        else:
+            raise ValueError(f"unsupported summary type: {summary_type}")
+
+    def compute(self) -> Dict[str, torch.Tensor]:
+        summary_type = self._summary_cfg.WhichOneof("summary")
+        result: Dict[str, torch.Tensor] = {}
+        if summary_type == "pcoc":
+            mean_pred = self._state.compute_mean()
+            mean_label = self._label_state.compute_mean()
+            if mean_pred is None or mean_label is None:
+                return result
+            epsilon = self._summary_cfg.pcoc.epsilon
+            result["summary/predicted_ctr"] = mean_pred
+            result["summary/observed_ctr"] = mean_label
+            result["summary/pcoc"] = mean_pred / (mean_label + epsilon)
+        elif summary_type == "scalars":
+            mean_pred = self._state.compute_mean()
+            if mean_pred is not None:
+                name = self._summary_cfg.scalars.name
+                result[f"summary/{name}"] = mean_pred
+        self.reset()
+        return result
+
+    def reset(self) -> None:
+        self._state.reset()
+        self._label_state.reset()
+
+
+def build_train_summary_modules(
+    summaries_set: List[ModelSummaries],
+) -> nn.ModuleList:
+    return nn.ModuleList([TrainSummaryModule(cfg) for cfg in summaries_set])