torch-rechub 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

torch_rechub/utils/model_utils.py

@@ -0,0 +1,233 @@
+"""Common model utility functions for Torch-RecHub.
+
+This module provides shared utilities for model introspection and input generation,
+used by both ONNX export and visualization features.
+
+Examples
+--------
+>>> from torch_rechub.utils.model_utils import extract_feature_info, generate_dummy_input
+>>> feature_info = extract_feature_info(model)
+>>> dummy_input = generate_dummy_input(feature_info['features'], batch_size=2)
+"""
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+# Import feature types for type checking
+try:
+    from ..basic.features import DenseFeature, SequenceFeature, SparseFeature
+except ImportError:
+    # Fallback for standalone usage
+    SparseFeature = None
+    DenseFeature = None
+    SequenceFeature = None
+
+
+def extract_feature_info(model: nn.Module) -> Dict[str, Any]:
+    """Extract feature information from a torch-rechub model using reflection.
+
+    This function inspects model attributes to find feature lists without
+    modifying the model code. Supports various model architectures.
+
+    Parameters
+    ----------
+    model : nn.Module
+        The recommendation model to inspect.
+
+    Returns
+    -------
+    dict
+        Dictionary containing:
+        - 'features': List of unique Feature objects
+        - 'input_names': List of feature names in order
+        - 'input_types': Dict mapping feature name to feature type
+        - 'user_features': List of user-side features (for dual-tower models)
+        - 'item_features': List of item-side features (for dual-tower models)
+
+    Examples
+    --------
+    >>> from torch_rechub.models.ranking import DeepFM
+    >>> model = DeepFM(deep_features, fm_features, mlp_params)
+    >>> info = extract_feature_info(model)
+    >>> print(info['input_names'])  # ['user_id', 'item_id', ...]
+    """
+    # Common feature attribute names across different model types
+    feature_attrs = [
+        'features',           # MMOE, DCN, etc.
+        'deep_features',      # DeepFM, WideDeep
+        'fm_features',        # DeepFM
+        'wide_features',      # WideDeep
+        'linear_features',    # DeepFFM
+        'cross_features',     # DeepFFM
+        'user_features',      # DSSM, YoutubeDNN, MIND
+        'item_features',      # DSSM, YoutubeDNN, MIND
+        'history_features',   # DIN, MIND
+        'target_features',    # DIN
+        'neg_item_feature',   # YoutubeDNN, MIND
+    ]
+
+    all_features = []
+    user_features = []
+    item_features = []
+
+    for attr in feature_attrs:
+        if hasattr(model, attr):
+            feat_list = getattr(model, attr)
+            if isinstance(feat_list, list) and len(feat_list) > 0:
+                all_features.extend(feat_list)
+                # Track user/item features for dual-tower models
+                if 'user' in attr or 'history' in attr:
+                    user_features.extend(feat_list)
+                elif 'item' in attr:
+                    item_features.extend(feat_list)
+
+    # Deduplicate features by name while preserving order
+    seen = set()
+    unique_features = []
+    for f in all_features:
+        if hasattr(f, 'name') and f.name not in seen:
+            seen.add(f.name)
+            unique_features.append(f)
+
+    # Deduplicate user/item features
+    seen_user = set()
+    unique_user = [f for f in user_features if hasattr(f, 'name') and f.name not in seen_user and not seen_user.add(f.name)]
+    seen_item = set()
+    unique_item = [f for f in item_features if hasattr(f, 'name') and f.name not in seen_item and not seen_item.add(f.name)]
+
+    # Build input names and types
+    input_names = [f.name for f in unique_features if hasattr(f, 'name')]
+    input_types = {f.name: type(f).__name__ for f in unique_features if hasattr(f, 'name')}
+
+    return {
+        'features': unique_features,
+        'input_names': input_names,
+        'input_types': input_types,
+        'user_features': unique_user,
+        'item_features': unique_item,
+    }
+
+
+def generate_dummy_input(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Tuple[torch.Tensor, ...]:
+    """Generate dummy input tensors based on feature definitions.
+
+    Parameters
+    ----------
+    features : list
+        List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
+    batch_size : int, default=2
+        Batch size for dummy input.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature.
+    device : str, default='cpu'
+        Device to create tensors on.
+
+    Returns
+    -------
+    tuple of Tensor
+        Tuple of tensors in the order of input features.
+
+    Examples
+    --------
+    >>> features = [SparseFeature("user_id", 1000), SequenceFeature("hist", 500)]
+    >>> dummy = generate_dummy_input(features, batch_size=4)
+    >>> # Returns (user_id_tensor[4], hist_tensor[4, 10])
+    """
+    # Dynamic import to handle feature types
+    from ..basic.features import DenseFeature, SequenceFeature, SparseFeature
+
+    inputs = []
+    for feat in features:
+        if isinstance(feat, SequenceFeature):
+            # Sequence features have shape [batch_size, seq_length]
+            tensor = torch.randint(0, feat.vocab_size, (batch_size, seq_length), device=device)
+        elif isinstance(feat, SparseFeature):
+            # Sparse features have shape [batch_size]
+            tensor = torch.randint(0, feat.vocab_size, (batch_size,), device=device)
+        elif isinstance(feat, DenseFeature):
+            # Dense features always have shape [batch_size, embed_dim]
+            tensor = torch.randn(batch_size, feat.embed_dim, device=device)
+        else:
+            raise TypeError(f"Unsupported feature type: {type(feat)}")
+        inputs.append(tensor)
+    return tuple(inputs)
+
+
+def generate_dummy_input_dict(features: List[Any], batch_size: int = 2, seq_length: int = 10, device: str = 'cpu') -> Dict[str, torch.Tensor]:
+    """Generate dummy input dict based on feature definitions.
+
+    Similar to generate_dummy_input but returns a dict mapping feature names
+    to tensors. This is the expected input format for torch-rechub models.
+
+    Parameters
+    ----------
+    features : list
+        List of Feature objects (SparseFeature, DenseFeature, SequenceFeature).
+    batch_size : int, default=2
+        Batch size for dummy input.
+    seq_length : int, default=10
+        Sequence length for SequenceFeature.
+    device : str, default='cpu'
+        Device to create tensors on.
+
+    Returns
+    -------
+    dict
+        Dict mapping feature names to tensors.
+
+    Examples
+    --------
+    >>> features = [SparseFeature("user_id", 1000)]
+    >>> dummy = generate_dummy_input_dict(features, batch_size=4)
+    >>> # Returns {"user_id": tensor[4]}
+    """
+    dummy_tuple = generate_dummy_input(features, batch_size, seq_length, device)
+    input_names = [f.name for f in features if hasattr(f, 'name')]
+    return {name: tensor for name, tensor in zip(input_names, dummy_tuple)}
+
+
+def generate_dynamic_axes(input_names: List[str], output_names: Optional[List[str]] = None, batch_dim: int = 0, include_seq_dim: bool = True, seq_features: Optional[List[str]] = None) -> Dict[str, Dict[int, str]]:
+    """Generate dynamic axes configuration for ONNX export.
+
+    Parameters
+    ----------
+    input_names : list of str
+        List of input tensor names.
+    output_names : list of str, optional
+        List of output tensor names. Default is ["output"].
+    batch_dim : int, default=0
+        Dimension index for batch size.
+    include_seq_dim : bool, default=True
+        Whether to include sequence dimension as dynamic.
+    seq_features : list of str, optional
+        List of feature names that are sequences.
+
+    Returns
+    -------
+    dict
+        Dynamic axes dict for torch.onnx.export.
+
+    Examples
+    --------
+    >>> axes = generate_dynamic_axes(["user_id", "item_id"], seq_features=["hist"])
+    >>> # Returns {"user_id": {0: "batch_size"}, "item_id": {0: "batch_size"}, ...}
+    """
+    if output_names is None:
+        output_names = ["output"]
+
+    dynamic_axes = {}
+
+    # Input axes
+    for name in input_names:
+        dynamic_axes[name] = {batch_dim: "batch_size"}
+        # Add sequence dimension for sequence features
+        if include_seq_dim and seq_features and name in seq_features:
+            dynamic_axes[name][1] = "seq_length"
+
+    # Output axes
+    for name in output_names:
+        dynamic_axes[name] = {batch_dim: "batch_size"}
+
+    return dynamic_axes

torch_rechub/utils/mtl.py

@@ -1,126 +1,136 @@
-import torch
-from torch.optim.optimizer import Optimizer
-from ..models.multi_task import MMOE, SharedBottom, PLE, AITM
-
-
-def shared_task_layers(model):
-    """get shared layers and task layers in multi-task model
-    Authors: Qida Dong, dongjidan@126.com
-
-    Args:
-        model (torch.nn.Module): only support `[MMOE, SharedBottom, PLE, AITM]`
-
-    Returns:
-        list[torch.nn.parameter]: parameters split to shared list and task list.
-    """
-    shared_layers = list(model.embedding.parameters())
-    task_layers = None
-    if isinstance(model, SharedBottom):
-        shared_layers += list(model.bottom_mlp.parameters())
-        task_layers = list(model.towers.parameters()) + list(model.predict_layers.parameters())
-    elif isinstance(model, MMOE):
-        shared_layers += list(model.experts.parameters())
-        task_layers = list(model.towers.parameters()) + list(model.predict_layers.parameters())
-        task_layers += list(model.gates.parameters())
-    elif isinstance(model, PLE):
-        shared_layers += list(model.cgc_layers.parameters())
-        task_layers = list(model.towers.parameters()) + list(model.predict_layers.parameters())
-    elif isinstance(model, AITM):
-        shared_layers += list(model.bottoms.parameters())
-        task_layers = list(model.info_gates.parameters()) + list(model.towers.parameters()) + list(
-            model.aits.parameters())
-    else:
-        raise ValueError(f'this model {model} is not suitable for MetaBalance Optimizer')
-    return shared_layers, task_layers
-
-
-class MetaBalance(Optimizer):
-    """MetaBalance Optimizer
-    This method is used to scale the gradient and balance the gradient of each task.
-    Authors: Qida Dong, dongjidan@126.com
-
-    Args:
-        parameters (list): the parameters of model
-        relax_factor (float, optional): the relax factor of gradient scaling (default: 0.7)
-        beta (float, optional): the coefficient of moving average (default: 0.9)
-    """
-
-    def __init__(self, parameters, relax_factor=0.7, beta=0.9):
-
-        if relax_factor < 0. or relax_factor >= 1.:
-            raise ValueError(f'Invalid relax_factor: {relax_factor}, it should be 0. <= relax_factor < 1.')
-        if beta < 0. or beta >= 1.:
-            raise ValueError(f'Invalid beta: {beta}, it should be 0. <= beta < 1.')
-        rel_beta_dict = {'relax_factor': relax_factor, 'beta': beta}
-        super(MetaBalance, self).__init__(parameters, rel_beta_dict)
-
-    @torch.no_grad()
-    def step(self, losses):
-        for idx, loss in enumerate(losses):
-            loss.backward(retain_graph=True)
-            for group in self.param_groups:
-                for gp in group['params']:
-                    if gp.grad is None:
-                        # print('breaking')
-                        break
-                    if gp.grad.is_sparse:
-                        raise RuntimeError('MetaBalance does not support sparse gradients')
-                    # store the result of moving average
-                    state = self.state[gp]
-                    if len(state) == 0:
-                        for i in range(len(losses)):
-                            if i == 0:
-                                gp.norms = [0]
-                            else:
-                                gp.norms.append(0)
-                    # calculate the moving average
-                    beta = group['beta']
-                    gp.norms[idx] = gp.norms[idx] * beta + (1 - beta) * torch.norm(gp.grad)
-                    # scale the auxiliary gradient
-                    relax_factor = group['relax_factor']
-                    gp.grad = gp.grad * gp.norms[0] / (gp.norms[idx] + 1e-5) * relax_factor + gp.grad * (1. -
-                                                                                                         relax_factor)
-                    # store the gradient of each auxiliary task in state
-                    if idx == 0:
-                        state['sum_gradient'] = torch.zeros_like(gp.data)
-                        state['sum_gradient'] += gp.grad
-                    else:
-                        state['sum_gradient'] += gp.grad
-
-                    if gp.grad is not None:
-                        gp.grad.detach_()
-                        gp.grad.zero_()
-                    if idx == len(losses) - 1:
-                        gp.grad = state['sum_gradient']
-
-
-def gradnorm(loss_list, loss_weight, share_layer, initial_task_loss, alpha):
-    loss = 0
-    for loss_i, w_i in zip(loss_list, loss_weight):
-        loss += loss_i * w_i
-    loss.backward(retain_graph=True)
-    # set the gradients of w_i(t) to zero because these gradients have to be updated using the GradNorm loss
-    for w_i in loss_weight:
-        w_i.grad.data = w_i.grad.data * 0.0
-    # get the gradient norms for each of the tasks
-    # G^{(i)}_w(t)
-    norms, loss_ratio = [], []
-    for i in range(len(loss_list)):
-        # get the gradient of this task loss with respect to the shared parameters
-        gygw = torch.autograd.grad(loss_list[i], share_layer, retain_graph=True)
-        # compute the norm
-        norms.append(torch.norm(torch.mul(loss_weight[i], gygw[0])))
-        # compute the inverse training rate r_i(t)
-        loss_ratio.append(loss_list[i].item() / initial_task_loss[i])
-    norms = torch.stack(norms)
-    mean_norm = torch.mean(norms.detach())
-    mean_loss_ratio = sum(loss_ratio) / len(loss_ratio)
-    # compute the GradNorm loss
-    # this term has to remain constant
-    constant_term = mean_norm * (mean_loss_ratio**alpha)
-    grad_norm_loss = torch.sum(torch.abs(norms - constant_term))
-    #print('GradNorm loss {}'.format(grad_norm_loss))
-
-    # compute the gradient for the weights
-    for w_i in loss_weight:
-        w_i.grad = torch.autograd.grad(grad_norm_loss, w_i, retain_graph=True)[0]
+import torch
+from torch.optim.optimizer import Optimizer
+
+from ..models.multi_task import AITM, MMOE, PLE, SharedBottom
+
+
+def shared_task_layers(model):
+    """get shared layers and task layers in multi-task model
+    Authors: Qida Dong, dongjidan@126.com
+
+    Args:
+        model (torch.nn.Module): only support `[MMOE, SharedBottom, PLE, AITM]`
+
+    Returns:
+        list[torch.nn.parameter]: parameters split to shared list and task list.
+    """
+    shared_layers = list(model.embedding.parameters())
+    task_layers = None
+    if isinstance(model, SharedBottom):
+        shared_layers += list(model.bottom_mlp.parameters())
+        task_layers = list(model.towers.parameters()) + \
+            list(model.predict_layers.parameters())
+    elif isinstance(model, MMOE):
+        shared_layers += list(model.experts.parameters())
+        task_layers = list(model.towers.parameters()) + \
+            list(model.predict_layers.parameters())
+        task_layers += list(model.gates.parameters())
+    elif isinstance(model, PLE):
+        shared_layers += list(model.cgc_layers.parameters())
+        task_layers = list(model.towers.parameters()) + \
+            list(model.predict_layers.parameters())
+    elif isinstance(model, AITM):
+        shared_layers += list(model.bottoms.parameters())
+        task_layers = list(model.info_gates.parameters()) + list(model.towers.parameters()) + list(model.aits.parameters())
+    else:
+        raise ValueError(f'this model {model} is not suitable for MetaBalance Optimizer')
+    return shared_layers, task_layers
+
+
+class MetaBalance(Optimizer):
+    """MetaBalance Optimizer
+    This method is used to scale the gradient and balance the gradient of each task.
+    Authors: Qida Dong, dongjidan@126.com
+
+    Args:
+        parameters (list): the parameters of model
+        relax_factor (float, optional): the relax factor of gradient scaling (default: 0.7)
+        beta (float, optional): the coefficient of moving average (default: 0.9)
+    """
+
+    def __init__(self, parameters, relax_factor=0.7, beta=0.9):
+
+        if relax_factor < 0. or relax_factor >= 1.:
+            raise ValueError(f'Invalid relax_factor: {relax_factor}, it should be 0. <= relax_factor < 1.')
+        if beta < 0. or beta >= 1.:
+            raise ValueError(f'Invalid beta: {beta}, it should be 0. <= beta < 1.')
+        rel_beta_dict = {'relax_factor': relax_factor, 'beta': beta}
+        super(MetaBalance, self).__init__(parameters, rel_beta_dict)
+
+    @torch.no_grad()
+    def step(self, losses):
+        for idx, loss in enumerate(losses):
+            loss.backward(retain_graph=True)
+            for group in self.param_groups:
+                for gp in group['params']:
+                    if gp.grad is None:
+                        # print('breaking')
+                        break
+                    if gp.grad.is_sparse:
+                        raise RuntimeError('MetaBalance does not support sparse gradients')
+# store the result of moving average
+                    state = self.state[gp]
+                    if len(state) == 0:
+                        for i in range(len(losses)):
+                            if i == 0:
+                                gp.norms = [0]
+                            else:
+                                gp.norms.append(0)
+
+
+# calculate the moving average
+                    beta = group['beta']
+                    gp.norms[idx] = gp.norms[idx] * beta + \
+                        (1 - beta) * torch.norm(gp.grad)
+                    # scale the auxiliary gradient
+                    relax_factor = group['relax_factor']
+                    gp.grad = gp.grad * \
+                        gp.norms[0] / (gp.norms[idx] + 1e-5) * relax_factor + gp.grad * (1. - relax_factor)
+                    # store the gradient of each auxiliary task in state
+                    if idx == 0:
+                        state['sum_gradient'] = torch.zeros_like(gp.data)
+                        state['sum_gradient'] += gp.grad
+                    else:
+                        state['sum_gradient'] += gp.grad
+
+                    if gp.grad is not None:
+                        gp.grad.detach_()
+                        gp.grad.zero_()
+                    if idx == len(losses) - 1:
+                        gp.grad = state['sum_gradient']
+
+
+def gradnorm(loss_list, loss_weight, share_layer, initial_task_loss, alpha):
+    loss = 0
+    for loss_i, w_i in zip(loss_list, loss_weight):
+        loss += loss_i * w_i
+    loss.backward(retain_graph=True)
+    # set the gradients of w_i(t) to zero because these gradients have to be
+    # updated using the GradNorm loss
+    for w_i in loss_weight:
+        w_i.grad.data = w_i.grad.data * 0.0
+
+
+# get the gradient norms for each of the tasks
+# G^{(i)}_w(t)
+    norms, loss_ratio = [], []
+    for i in range(len(loss_list)):
+        # get the gradient of this task loss with respect to the shared
+        # parameters
+        gygw = torch.autograd.grad(loss_list[i], share_layer, retain_graph=True)
+        # compute the norm
+        norms.append(torch.norm(torch.mul(loss_weight[i], gygw[0])))
+        # compute the inverse training rate r_i(t)
+        loss_ratio.append(loss_list[i].item() / initial_task_loss[i])
+    norms = torch.stack(norms)
+    mean_norm = torch.mean(norms.detach())
+    mean_loss_ratio = sum(loss_ratio) / len(loss_ratio)
+    # compute the GradNorm loss
+    # this term has to remain constant
+    constant_term = mean_norm * (mean_loss_ratio**alpha)
+    grad_norm_loss = torch.sum(torch.abs(norms - constant_term))
+    # print('GradNorm loss {}'.format(grad_norm_loss))
+
+    # compute the gradient for the weights
+    for w_i in loss_weight:
+        w_i.grad = torch.autograd.grad(grad_norm_loss, w_i, retain_graph=True)[0]