nextrec 0.4.32__py3-none-any.whl → 0.4.34__py3-none-any.whl

nextrec/loss/grad_norm.py

@@ -2,12 +2,40 @@
 GradNorm loss weighting for multi-task learning.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 24/12/2025
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou,zyaztec@gmail.com
 
 Reference:
 Chen, Zhao, et al. "GradNorm: Gradient Normalization for Adaptive Loss Balancing
 in Deep Multitask Networks." ICML 2018.
+
+pseudocode:
+---
+Initialize w_i = 1
+Record L_i(0)
+
+for each step:
+    1. Forward: compute each task loss L_i
+    2. Compute G_i = ||∇_W (w_i * L_i)||
+    3. Compute r_i = (L_i / L_i(0)) / mean(...)
+    4. Compute target: Ĝ_i = mean(G) * r_i^α
+    5. L_grad = sum |G_i - Ĝ_i|
+    6. Update w_i using ∇ L_grad
+    7. Backprop with sum_i (w_i * L_i) to update model
+
+伪代码：
+---
+初始化 w_i = 1
+记录 L_i(0)
+
+for each step:
+    1. 前向算各 task loss: L_i
+    2. 计算 G_i = ||∇_W (w_i * L_i)||
+    3. 计算 r_i = (L_i / L_i(0)) / mean(...)
+    4. 计算 target: Ĝ_i = mean(G) * r_i^α
+    5. L_grad = sum |G_i - Ĝ_i|
+    6. 对 w_i 用 ∇L_grad 更新
+    7. 用 ∑ w_i * L_i 反传更新模型
 """
 
 from __future__ import annotations
@@ -15,6 +43,7 @@ from __future__ import annotations
 from typing import Iterable
 
 import torch
+import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 
@@ -23,7 +52,15 @@ def get_grad_norm_shared_params(
     model,
     shared_modules=None,
 ):
+    """
+    Get shared parameters for GradNorm.
+
+    Args:
+        model: A pytorch model instance containing grad_norm_shared_modules attribute.
+        shared_modules: Optional list of module names to consider as shared.
+    """
     if not shared_modules:
+        # If no specific shared modules are provided, consider all parameters as shared
         return [p for p in model.parameters() if p.requires_grad]
     shared_params = []
     seen = set()
@@ -35,26 +72,10 @@ def get_grad_norm_shared_params(
             if param.requires_grad and id(param) not in seen:
                 shared_params.append(param)
                 seen.add(id(param))
-    if not shared_params:
-        return [p for p in model.parameters() if p.requires_grad]
     return shared_params
 
 
 class GradNormLossWeighting:
-    """
-    Adaptive multi-task loss weighting with GradNorm.
-
-    Args:
-        nums_task: Number of tasks.
-        alpha: GradNorm balancing strength.
-        lr: Learning rate for the weight optimizer.
-        init_weights: Optional initial weights per task.
-        device: Torch device for weights.
-        ema_decay: Optional EMA decay for smoothing loss ratios.
-        init_ema_steps: Number of steps to build EMA for initial losses.
-        init_ema_decay: EMA decay for initial losses when init_ema_steps > 0.
-        eps: Small value for numerical stability.
-    """
 
     def __init__(
         self,
@@ -63,58 +84,43 @@ class GradNormLossWeighting:
         lr: float = 0.025,
         init_weights: Iterable[float] | None = None,
         device: torch.device | str | None = None,
-        ema_decay: float | None = None,
-        init_ema_steps: int = 0,
-        init_ema_decay: float = 0.9,
         eps: float = 1e-8,
     ) -> None:
+        """
+        Adaptive multi-task loss weighting with GradNorm.
+
+        Args:
+            nums_task: Number of tasks.
+            alpha: GradNorm balancing strength.
+            lr: Learning rate for the weight optimizer.
+            init_weights: Optional initial weights per task.
+            device: Torch device for weights.
+            eps: Small value for numerical stability.
+
+
+
+        """
+
         if nums_task <= 1:
             raise ValueError("GradNorm requires nums_task > 1.")
+
         self.nums_task = nums_task
         self.alpha = alpha
         self.eps = eps
-        if ema_decay is not None:
-            ema_decay = ema_decay
-            if ema_decay < 0.0 or ema_decay >= 1.0:
-                raise ValueError("ema_decay must be in [0.0, 1.0).")
-        self.ema_decay = ema_decay
-        self.init_ema_steps = init_ema_steps
-        if self.init_ema_steps < 0:
-            raise ValueError("init_ema_steps must be >= 0.")
-        self.init_ema_decay = init_ema_decay
-        if self.init_ema_decay < 0.0 or self.init_ema_decay >= 1.0:
-            raise ValueError("init_ema_decay must be in [0.0, 1.0).")
-        self.init_ema_count = 0
 
         if init_weights is None:
             weights = torch.ones(self.nums_task, dtype=torch.float32)
         else:
             weights = torch.tensor(list(init_weights), dtype=torch.float32)
-            if weights.numel() != self.nums_task:
-                raise ValueError(
-                    "init_weights length must match nums_task for GradNorm."
-                )
+
         if device is not None:
             weights = weights.to(device)
         self.weights = nn.Parameter(weights)
         self.optimizer = torch.optim.Adam([self.weights], lr=float(lr))
 
         self.initial_losses = None
-        self.initial_losses_ema = None
-        self.loss_ema = None
         self.pending_grad = None
 
-    def to(self, device):
-        device = torch.device(device)
-        self.weights.data = self.weights.data.to(device)
-        if self.initial_losses is not None:
-            self.initial_losses = self.initial_losses.to(device)
-        if self.initial_losses_ema is not None:
-            self.initial_losses_ema = self.initial_losses_ema.to(device)
-        if self.loss_ema is not None:
-            self.loss_ema = self.loss_ema.to(device)
-        return self
-
     def compute_weighted_loss(
         self,
         task_losses: list[torch.Tensor],
@@ -122,6 +128,8 @@ class GradNormLossWeighting:
     ) -> torch.Tensor:
         """
         Return weighted total loss and update task weights with GradNorm.
+
+        BaseModel will use this method to compute the weighted loss when self.grad_norm is enabled.
         """
         if len(task_losses) != self.nums_task:
             raise ValueError(
@@ -136,19 +144,7 @@ class GradNormLossWeighting:
                 [loss.item() for loss in task_losses], device=self.weights.device
             )
             if self.initial_losses is None:
-                if self.init_ema_steps > 0:
-                    if self.initial_losses_ema is None:
-                        self.initial_losses_ema = loss_values
-                    else:
-                        self.initial_losses_ema = (
-                            self.init_ema_decay * self.initial_losses_ema
-                            + (1.0 - self.init_ema_decay) * loss_values
-                        )
-                    self.init_ema_count += 1
-                    if self.init_ema_count >= self.init_ema_steps:
-                        self.initial_losses = self.initial_losses_ema.clone()
-                else:
-                    self.initial_losses = loss_values
+                self.initial_losses = loss_values.clone()
 
         weights_detached = self.weights.detach()
         weighted_losses = [
@@ -157,25 +153,14 @@ class GradNormLossWeighting:
         total_loss = torch.stack(weighted_losses).sum()
 
         grad_norms = self.compute_grad_norms(task_losses, shared_params)
+
+        # compute inverse training rate, inv rate = loss_ratio / mean(loss_ratio)
         with torch.no_grad():
-            if self.ema_decay is not None:
-                if self.loss_ema is None:
-                    self.loss_ema = loss_values
-                else:
-                    self.loss_ema = (
-                        self.ema_decay * self.loss_ema
-                        + (1.0 - self.ema_decay) * loss_values
-                    )
-                ratio_source = self.loss_ema
-            else:
-                ratio_source = loss_values
             if self.initial_losses is not None:
                 base_initial = self.initial_losses
-            elif self.initial_losses_ema is not None:
-                base_initial = self.initial_losses_ema
             else:
                 base_initial = loss_values
-            loss_ratios = ratio_source / (base_initial + self.eps)
+            loss_ratios = loss_values / (base_initial + self.eps)
             inv_rate = loss_ratios / (loss_ratios.mean() + self.eps)
             target = grad_norms.mean() * (inv_rate**self.alpha)
 
@@ -187,6 +172,7 @@ class GradNormLossWeighting:
 
     def compute_grad_norms(self, task_losses, shared_params):
         grad_norms = []
+        # compute gradient norms for each task, gradient norms = sqrt(sum(grad^2))
         for i, task_loss in enumerate(task_losses):
             grads = torch.autograd.grad(
                 self.weights[i] * task_loss,
@@ -230,3 +216,19 @@ class GradNormLossWeighting:
             self.weights.copy_(w)
 
         self.pending_grad = None
+
+    def sync(self) -> None:
+        """
+        Synchronize GradNorm buffers across DDP ranks.
+
+        - pending_grad: averaged so all ranks update weights consistently
+        - initial_losses: averaged so the baseline loss is consistent
+        """
+
+        world_size = dist.get_world_size()
+        if self.pending_grad is not None:
+            dist.all_reduce(self.pending_grad, op=dist.ReduceOp.SUM)
+            self.pending_grad /= world_size
+        if self.initial_losses is not None:
+            dist.all_reduce(self.initial_losses, op=dist.ReduceOp.SUM)
+            self.initial_losses /= world_size

nextrec/models/multi_task/ple.py

@@ -54,6 +54,7 @@ from nextrec.basic.model import BaseModel
 from nextrec.utils.model import get_mlp_output_dim
 from nextrec.utils.types import TaskTypeInput
 
+
 class CGCLayer(nn.Module):
     """
     CGC (Customized Gate Control) block used by PLE.

nextrec/models/multi_task/share_bottom.py

@@ -45,6 +45,7 @@ from nextrec.basic.heads import TaskHead
 from nextrec.basic.model import BaseModel
 from nextrec.utils.types import TaskTypeInput
 
+
 class ShareBottom(BaseModel):
     @property
     def model_name(self):

nextrec/models/tree_base/base.py

@@ -29,7 +29,7 @@ from nextrec.data.dataloader import RecDataLoader
 from nextrec.data.data_processing import get_column_data
 from nextrec.utils.console import display_metrics_table
 from nextrec.utils.data import FILE_FORMAT_CONFIG, check_streaming_support
-from nextrec.utils.feature import to_list
+from nextrec.utils.torch_utils import to_list
 from nextrec.utils.torch_utils import to_numpy
 
 

nextrec/utils/__init__.py

@@ -36,7 +36,7 @@ from .data import (
     resolve_file_paths,
 )
 from .embedding import get_auto_embedding_dim
-from .feature import to_list
+from .torch_utils import as_float, to_list
 from .model import (
     compute_pair_scores,
     get_mlp_output_dim,
@@ -90,6 +90,7 @@ __all__ = [
     "normalize_task_loss",
     # Feature utilities
     "to_list",
+    "as_float",
     # Config utilities
     "resolve_path",
     "safe_value",

nextrec/utils/config.py

@@ -21,7 +21,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Tuple
 import pandas as pd
 import torch
 
-from nextrec.utils.feature import to_list
+from nextrec.utils.torch_utils import to_list
 
 if TYPE_CHECKING:
     from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature

nextrec/utils/console.py

@@ -36,7 +36,7 @@ from rich.progress import (
 from rich.table import Table
 from rich.text import Text
 
-from nextrec.utils.feature import as_float, to_list
+from nextrec.utils.torch_utils import as_float, to_list
 
 T = TypeVar("T")
 

nextrec/utils/torch_utils.py

@@ -5,14 +5,15 @@ This module groups device setup, distributed helpers, optimizers/schedulers,
 initialization, and tensor helpers.
 
 Date: create on 27/10/2025
-Checkpoint: edit on 27/12/2025
+Checkpoint: edit on 22/01/2026
 Author: Yang Zhou, zyaztec@gmail.com
 """
 
 from __future__ import annotations
 
 import logging
-from typing import Any, Dict, Iterable, Literal
+import numbers
+from typing import Any, Dict, Iterable
 
 import numpy as np
 import torch
@@ -22,7 +23,55 @@ from torch.utils.data import DataLoader, IterableDataset
 from torch.utils.data.distributed import DistributedSampler
 
 from nextrec.basic.loggers import colorize
-from nextrec.utils.types import OptimizerName, SchedulerName
+from nextrec.utils.types import (
+    EmbeddingInitType,
+    InitializerActivationType,
+    OptimizerName,
+    SchedulerName,
+)
+
+
+def to_list(value: str | list[str] | None) -> list[str]:
+    if value is None:
+        return []
+    if isinstance(value, str):
+        return [value]
+    return list(value)
+
+
+def as_float(value: Any) -> float | None:
+    if isinstance(value, numbers.Number):
+        return float(value)
+    if hasattr(value, "item"):
+        try:
+            return float(value.item())
+        except Exception:
+            return None
+    return None
+
+
+def to_numpy(values: Any) -> np.ndarray:
+    if isinstance(values, torch.Tensor):
+        return values.detach().cpu().numpy()
+    return np.asarray(values)
+
+
+def to_tensor(
+    value: Any, dtype: torch.dtype, device: torch.device | str | None = None
+) -> torch.Tensor:
+    if value is None:
+        raise ValueError("[Tensor Utils Error] Cannot convert None to tensor.")
+    tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value)
+    if tensor.dtype != dtype:
+        tensor = tensor.to(dtype=dtype)
+
+    if device is not None:
+        target_device = (
+            device if isinstance(device, torch.device) else torch.device(device)
+        )
+        if tensor.device != target_device:
+            tensor = tensor.to(target_device)
+    return tensor
 
 
 def resolve_nonlinearity(activation: str) -> str:
@@ -56,30 +105,8 @@ def resolve_gain(activation: str, param: Dict[str, Any]) -> float:
 
 
 def get_initializer(
-    init_type: Literal[
-        "xavier_uniform",
-        "xavier_normal",
-        "kaiming_uniform",
-        "kaiming_normal",
-        "orthogonal",
-        "normal",
-        "uniform",
-    ] = "normal",
-    activation: Literal[
-        "linear",
-        "conv1d",
-        "conv2d",
-        "conv3d",
-        "conv_transpose1d",
-        "conv_transpose2d",
-        "conv_transpose3d",
-        "sigmoid",
-        "tanh",
-        "relu",
-        "leaky_relu",
-        "selu",
-        "gelu",
-    ] = "linear",
+    init_type: EmbeddingInitType = "normal",
+    activation: InitializerActivationType = "linear",
     param: Dict[str, Any] | None = None,
 ):
     param = param or {}
@@ -108,7 +135,7 @@ def get_initializer(
         elif init_type == "uniform":
             nn.init.uniform_(tensor, a=param.get("a", -0.05), b=param.get("b", 0.05))
         else:
-            raise ValueError(f"Unknown init_type: {init_type}")
+            raise ValueError(f"[Initializer Error] Unknown init_type: {init_type}")
         return tensor
 
     return initializer_fn
@@ -172,12 +199,14 @@ def get_optimizer(
         elif opt_name == "rmsprop":
             opt_class = torch.optim.RMSprop
         else:
-            raise NotImplementedError(f"Unsupported optimizer: {optimizer}")
+            raise NotImplementedError(
+                f"[Optimizer Error] Unsupported optimizer: {optimizer}"
+            )
         optimizer_fn = opt_class(params=params, **optimizer_params)
     elif isinstance(optimizer, torch.optim.Optimizer):
         optimizer_fn = optimizer
     else:
-        raise TypeError(f"Invalid optimizer type: {type(optimizer)}")
+        raise TypeError(f"[Optimizer Error] Invalid optimizer type: {type(optimizer)}")
     return optimizer_fn
 
 
@@ -203,7 +232,9 @@ def get_scheduler(
                 optimizer, **scheduler_params
             )
         else:
-            raise NotImplementedError(f"Unsupported scheduler: {scheduler}")
+            raise NotImplementedError(
+                f"[Scheduler Error] Unsupported scheduler: {scheduler}"
+            )
     elif isinstance(scheduler, type) and issubclass(
         scheduler,
         (torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.LRScheduler),
@@ -215,35 +246,11 @@ def get_scheduler(
     ):
         scheduler_fn = scheduler
     else:
-        raise TypeError(f"Invalid scheduler type: {type(scheduler)}")
+        raise TypeError(f"[Scheduler Error] Invalid scheduler type: {type(scheduler)}")
 
     return scheduler_fn
 
 
-def to_numpy(values: Any) -> np.ndarray:
-    if isinstance(values, torch.Tensor):
-        return values.detach().cpu().numpy()
-    return np.asarray(values)
-
-
-def to_tensor(
-    value: Any, dtype: torch.dtype, device: torch.device | str | None = None
-) -> torch.Tensor:
-    if value is None:
-        raise ValueError("[Tensor Utils Error] Cannot convert None to tensor.")
-    tensor = value if isinstance(value, torch.Tensor) else torch.as_tensor(value)
-    if tensor.dtype != dtype:
-        tensor = tensor.to(dtype=dtype)
-
-    if device is not None:
-        target_device = (
-            device if isinstance(device, torch.device) else torch.device(device)
-        )
-        if tensor.device != target_device:
-            tensor = tensor.to(target_device)
-    return tensor
-
-
 def init_process_group(
     distributed: bool, rank: int, world_size: int, device_id: int | None = None
 ) -> None:

nextrec/utils/types.py

@@ -64,6 +64,40 @@ TaskTypeName = Literal["binary", "regression"]
 
 TaskTypeInput = TaskTypeName | str
 
+EmbeddingInitType = Literal[
+    "normal",
+    "uniform",
+    "xavier_uniform",
+    "xavier_normal",
+    "kaiming_uniform",
+    "kaiming_normal",
+    "orthogonal",
+]
+
+SequenceCombinerType = Literal[
+    "mean",
+    "sum",
+    "concat",
+    "dot_attention",
+    "self_attention",
+]
+
+InitializerActivationType = Literal[
+    "linear",
+    "conv1d",
+    "conv2d",
+    "conv3d",
+    "conv_transpose1d",
+    "conv_transpose2d",
+    "conv_transpose3d",
+    "sigmoid",
+    "tanh",
+    "relu",
+    "leaky_relu",
+    "selu",
+    "gelu",
+]
+
 MetricsName = Literal[
     "auc",
     "gauc",
@@ -97,4 +131,13 @@ MetricsName = Literal[
     "mrr@5",
     "mrr@10",
     "mrr@20",
+    "topk_recall@5",
+    "topk_recall@10",
+    "topk_recall@20",
+    "topk_precision@5",
+    "topk_precision@10",
+    "topk_precision@20",
+    "lift@5",
+    "lift@10",
+    "lift@20",
 ]

{nextrec-0.4.32.dist-info → nextrec-0.4.34.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.4.32
+Version: 0.4.34
 Summary: A comprehensive recommendation library with match, ranking, and multi-task learning models
 Project-URL: Homepage, https://github.com/zerolovesea/NextRec
 Project-URL: Repository, https://github.com/zerolovesea/NextRec
@@ -69,7 +69,7 @@ Description-Content-Type: text/markdown
 ![Python](https://img.shields.io/badge/Python-3.10+-blue.svg)
 ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg)
 ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg)
-![Version](https://img.shields.io/badge/Version-0.4.32-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.4.34-orange.svg)
 [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec)
 
 中文文档 | [English Version](README_en.md)
@@ -254,11 +254,11 @@ nextrec --mode=predict --predict_config=path/to/predict_config.yaml
 
 预测结果固定保存到 `{checkpoint_path}/predictions/{name}.{save_data_format}`。
 
-> 截止当前版本0.4.32，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
+> 截止当前版本0.4.34，NextRec CLI支持单机训练，分布式训练相关功能尚在开发中。
 
 ## 兼容平台
 
-当前最新版本为0.4.32，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
+当前最新版本为0.4.34，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本：
 
 | 平台 | 配置 | 
 |------|------|