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

torch_rechub/trainers/match_trainer.py

@@ -39,6 +39,7 @@ class MatchTrainer(object):
         device="cpu",
         gpus=None,
         model_path="./",
+        model_logger=None,
     ):
         self.model = model  # for uniform weights save method in one gpu or multi gpu
         if gpus is None:
@@ -73,10 +74,13 @@ class MatchTrainer(object):
         self.model_path = model_path
         # Initialize regularization loss
         self.reg_loss_fn = RegularizationLoss(**regularization_params)
+        self.model_logger = model_logger
 
     def train_one_epoch(self, data_loader, log_interval=10):
         self.model.train()
         total_loss = 0
+        epoch_loss = 0
+        batch_count = 0
         tk0 = tqdm.tqdm(data_loader, desc="train", smoothing=0, mininterval=1.0)
         for i, (x_dict, y) in enumerate(tk0):
             x_dict = {k: v.to(self.device) for k, v in x_dict.items()}  # tensor to GPU
@@ -114,14 +118,26 @@ class MatchTrainer(object):
             loss.backward()
             self.optimizer.step()
             total_loss += loss.item()
+            epoch_loss += loss.item()
+            batch_count += 1
             if (i + 1) % log_interval == 0:
                 tk0.set_postfix(loss=total_loss / log_interval)
                 total_loss = 0
 
+        # Return average epoch loss
+        return epoch_loss / batch_count if batch_count > 0 else 0
+
     def fit(self, train_dataloader, val_dataloader=None):
+        for logger in self._iter_loggers():
+            logger.log_hyperparams({'n_epoch': self.n_epoch, 'learning_rate': self.optimizer.param_groups[0]['lr'], 'loss_mode': self.mode})
+
         for epoch_i in range(self.n_epoch):
             print('epoch:', epoch_i)
-            self.train_one_epoch(train_dataloader)
+            train_loss = self.train_one_epoch(train_dataloader)
+
+            for logger in self._iter_loggers():
+                logger.log_metrics({'train/loss': train_loss, 'learning_rate': self.optimizer.param_groups[0]['lr']}, step=epoch_i)
+
             if self.scheduler is not None:
                 if epoch_i % self.scheduler.step_size == 0:
                     print("Current lr : {}".format(self.optimizer.state_dict()['param_groups'][0]['lr']))
@@ -130,12 +146,34 @@ class MatchTrainer(object):
             if val_dataloader:
                 auc = self.evaluate(self.model, val_dataloader)
                 print('epoch:', epoch_i, 'validation: auc:', auc)
+
+                for logger in self._iter_loggers():
+                    logger.log_metrics({'val/auc': auc}, step=epoch_i)
+
                 if self.early_stopper.stop_training(auc, self.model.state_dict()):
                     print(f'validation: best auc: {self.early_stopper.best_auc}')
                     self.model.load_state_dict(self.early_stopper.best_weights)
                     break
+
         torch.save(self.model.state_dict(), os.path.join(self.model_path, "model.pth"))  # save best auc model
 
+        for logger in self._iter_loggers():
+            logger.finish()
+
+    def _iter_loggers(self):
+        """Return logger instances as a list.
+
+        Returns
+        -------
+        list
+            Active logger instances. Empty when ``model_logger`` is ``None``.
+        """
+        if self.model_logger is None:
+            return []
+        if isinstance(self.model_logger, (list, tuple)):
+            return list(self.model_logger)
+        return [self.model_logger]
+
     def evaluate(self, model, data_loader):
         model.eval()
         targets, predicts = list(), list()
@@ -177,7 +215,7 @@ class MatchTrainer(object):
                 predicts.append(y_pred.data)
         return torch.cat(predicts, dim=0)
 
-    def export_onnx(self, output_path, mode=None, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False):
+    def export_onnx(self, output_path, mode=None, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False, onnx_export_kwargs=None):
         """Export the trained matching model to ONNX format.
 
         This method exports matching/retrieval models (e.g., DSSM, YoutubeDNN, MIND)
@@ -199,6 +237,7 @@ class MatchTrainer(object):
             device (str, optional): Device for export ('cpu', 'cuda', etc.).
                 If None, defaults to 'cpu' for maximum compatibility.
             verbose (bool): Print export details (default: False).
+            onnx_export_kwargs (dict, optional): Extra kwargs forwarded to ``torch.onnx.export``.
 
         Returns:
             bool: True if export succeeded, False otherwise.
@@ -232,7 +271,17 @@ class MatchTrainer(object):
 
         try:
             exporter = ONNXExporter(model, device=export_device)
-            return exporter.export(output_path=output_path, mode=mode, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+            return exporter.export(
+                output_path=output_path,
+                mode=mode,
+                dummy_input=dummy_input,
+                batch_size=batch_size,
+                seq_length=seq_length,
+                opset_version=opset_version,
+                dynamic_batch=dynamic_batch,
+                verbose=verbose,
+                onnx_export_kwargs=onnx_export_kwargs,
+            )
         finally:
             # Restore original mode
             if hasattr(model, 'mode'):

torch_rechub/trainers/mtl_trainer.py

@@ -47,6 +47,7 @@ class MTLTrainer(object):
         device="cpu",
         gpus=None,
         model_path="./",
+        model_logger=None,
     ):
         self.model = model
         if gpus is None:
@@ -104,6 +105,7 @@ class MTLTrainer(object):
         self.model_path = model_path
         # Initialize regularization loss
         self.reg_loss_fn = RegularizationLoss(**regularization_params)
+        self.model_logger = model_logger
 
     def train_one_epoch(self, data_loader):
         self.model.train()
@@ -163,21 +165,42 @@ class MTLTrainer(object):
     def fit(self, train_dataloader, val_dataloader, mode='base', seed=0):
         total_log = []
 
+        # Log hyperparameters once
+        for logger in self._iter_loggers():
+            logger.log_hyperparams({'n_epoch': self.n_epoch, 'learning_rate': self._current_lr(), 'adaptive_method': self.adaptive_method})
+
         for epoch_i in range(self.n_epoch):
             _log_per_epoch = self.train_one_epoch(train_dataloader)
 
+            # Collect metrics
+            logs = {f'train/task_{task_id}_loss': loss_val for task_id, loss_val in enumerate(_log_per_epoch)}
+            lr_value = self._current_lr()
+            if lr_value is not None:
+                logs['learning_rate'] = lr_value
+
             if self.scheduler is not None:
                 if epoch_i % self.scheduler.step_size == 0:
                     print("Current lr : {}".format(self.optimizer.state_dict()['param_groups'][0]['lr']))
                 self.scheduler.step()  # update lr in epoch level by scheduler
+
             scores = self.evaluate(self.model, val_dataloader)
             print('epoch:', epoch_i, 'validation scores: ', scores)
 
-            for score in scores:
+            for task_id, score in enumerate(scores):
+                logs[f'val/task_{task_id}_score'] = score
                 _log_per_epoch.append(score)
+            logs['auc'] = scores[self.earlystop_taskid]
+
+            if self.loss_weight:
+                for task_id, weight in enumerate(self.loss_weight):
+                    logs[f'loss_weight/task_{task_id}'] = weight.item()
 
             total_log.append(_log_per_epoch)
 
+            # Log metrics once per epoch
+            for logger in self._iter_loggers():
+                logger.log_metrics(logs, step=epoch_i)
+
             if self.early_stopper.stop_training(scores[self.earlystop_taskid], self.model.state_dict()):
                 print('validation best auc of main task %d: %.6f' % (self.earlystop_taskid, self.early_stopper.best_auc))
                 self.model.load_state_dict(self.early_stopper.best_weights)
@@ -185,8 +208,33 @@ class MTLTrainer(object):
 
         torch.save(self.model.state_dict(), os.path.join(self.model_path, "model_{}_{}.pth".format(mode, seed)))  # save best auc model
 
+        for logger in self._iter_loggers():
+            logger.finish()
+
         return total_log
 
+    def _iter_loggers(self):
+        """Return logger instances as a list.
+
+        Returns
+        -------
+        list
+            Active logger instances. Empty when ``model_logger`` is ``None``.
+        """
+        if self.model_logger is None:
+            return []
+        if isinstance(self.model_logger, (list, tuple)):
+            return list(self.model_logger)
+        return [self.model_logger]
+
+    def _current_lr(self):
+        """Fetch current learning rate regardless of adaptive method."""
+        if self.adaptive_method == "metabalance":
+            return self.share_optimizer.param_groups[0]['lr'] if hasattr(self, 'share_optimizer') else None
+        if hasattr(self, 'optimizer'):
+            return self.optimizer.param_groups[0]['lr']
+        return None
+
     def evaluate(self, model, data_loader):
         model.eval()
         targets, predicts = list(), list()
@@ -213,7 +261,7 @@ class MTLTrainer(object):
                 predicts.extend(y_preds.tolist())
         return predicts
 
-    def export_onnx(self, output_path, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False):
+    def export_onnx(self, output_path, dummy_input=None, batch_size=2, seq_length=10, opset_version=14, dynamic_batch=True, device=None, verbose=False, onnx_export_kwargs=None):
         """Export the trained multi-task model to ONNX format.
 
         This method exports multi-task learning models (e.g., MMOE, PLE, ESMM, SharedBottom)
@@ -235,6 +283,7 @@ class MTLTrainer(object):
             device (str, optional): Device for export ('cpu', 'cuda', etc.).
                 If None, defaults to 'cpu' for maximum compatibility.
             verbose (bool): Print export details (default: False).
+            onnx_export_kwargs (dict, optional): Extra kwargs forwarded to ``torch.onnx.export``.
 
         Returns:
             bool: True if export succeeded, False otherwise.
@@ -256,7 +305,16 @@ class MTLTrainer(object):
         export_device = device if device is not None else 'cpu'
 
         exporter = ONNXExporter(model, device=export_device)
-        return exporter.export(output_path=output_path, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+        return exporter.export(
+            output_path=output_path,
+            dummy_input=dummy_input,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            opset_version=opset_version,
+            dynamic_batch=dynamic_batch,
+            verbose=verbose,
+            onnx_export_kwargs=onnx_export_kwargs,
+        )
 
     def visualization(self, input_data=None, batch_size=2, seq_length=10, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
         """Visualize the model's computation graph.

torch_rechub/trainers/seq_trainer.py

@@ -46,7 +46,22 @@ class SeqTrainer(object):
         ... )
     """
 
-    def __init__(self, model, optimizer_fn=torch.optim.Adam, optimizer_params=None, scheduler_fn=None, scheduler_params=None, n_epoch=10, earlystop_patience=10, device='cpu', gpus=None, model_path='./', loss_type='cross_entropy', loss_params=None):
+    def __init__(
+        self,
+        model,
+        optimizer_fn=torch.optim.Adam,
+        optimizer_params=None,
+        scheduler_fn=None,
+        scheduler_params=None,
+        n_epoch=10,
+        earlystop_patience=10,
+        device='cpu',
+        gpus=None,
+        model_path='./',
+        loss_type='cross_entropy',
+        loss_params=None,
+        model_logger=None
+    ):
         self.model = model  # for uniform weights save method in one gpu or multi gpu
         if gpus is None:
             gpus = []
@@ -74,9 +89,11 @@ class SeqTrainer(object):
                 loss_params = {"ignore_index": 0}
             self.loss_fn = nn.CrossEntropyLoss(**loss_params)
 
+        self.loss_type = loss_type
         self.n_epoch = n_epoch
         self.early_stopper = EarlyStopper(patience=earlystop_patience)
         self.model_path = model_path
+        self.model_logger = model_logger
 
     def fit(self, train_dataloader, val_dataloader=None):
         """训练模型.
@@ -90,10 +107,18 @@ class SeqTrainer(object):
         """
         history = {'train_loss': [], 'val_loss': [], 'val_accuracy': []}
 
+        for logger in self._iter_loggers():
+            logger.log_hyperparams({'n_epoch': self.n_epoch, 'learning_rate': self.optimizer.param_groups[0]['lr'], 'loss_type': self.loss_type})
+
         for epoch_i in range(self.n_epoch):
             print('epoch:', epoch_i)
             # 训练阶段
-            self.train_one_epoch(train_dataloader)
+            train_loss = self.train_one_epoch(train_dataloader)
+            history['train_loss'].append(train_loss)
+
+            # Collect metrics
+            logs = {'train/loss': train_loss, 'learning_rate': self.optimizer.param_groups[0]['lr']}
+
             if self.scheduler is not None:
                 if epoch_i % self.scheduler.step_size == 0:
                     print("Current lr : {}".format(self.optimizer.state_dict()['param_groups'][0]['lr']))
@@ -105,6 +130,10 @@ class SeqTrainer(object):
                 history['val_loss'].append(val_loss)
                 history['val_accuracy'].append(val_accuracy)
 
+                logs['val/loss'] = val_loss
+                logs['val/accuracy'] = val_accuracy
+                logs['auc'] = val_accuracy  # For compatibility with EarlyStopper
+
                 print(f"epoch: {epoch_i}, validation: loss: {val_loss:.4f}, accuracy: {val_accuracy:.4f}")
 
                 # 早停
@@ -113,9 +142,30 @@ class SeqTrainer(object):
                     self.model.load_state_dict(self.early_stopper.best_weights)
                     break
 
+            for logger in self._iter_loggers():
+                logger.log_metrics(logs, step=epoch_i)
+
         torch.save(self.model.state_dict(), os.path.join(self.model_path, "model.pth"))  # save best model
+
+        for logger in self._iter_loggers():
+            logger.finish()
+
         return history
 
+    def _iter_loggers(self):
+        """Return logger instances as a list.
+
+        Returns
+        -------
+        list
+            Active logger instances. Empty when ``model_logger`` is ``None``.
+        """
+        if self.model_logger is None:
+            return []
+        if isinstance(self.model_logger, (list, tuple)):
+            return list(self.model_logger)
+        return [self.model_logger]
+
     def train_one_epoch(self, data_loader, log_interval=10):
         """Train the model for a single epoch.
 
@@ -128,6 +178,8 @@ class SeqTrainer(object):
         """
         self.model.train()
         total_loss = 0
+        epoch_loss = 0
+        batch_count = 0
         tk0 = tqdm.tqdm(data_loader, desc="train", smoothing=0, mininterval=1.0)
         for i, (seq_tokens, seq_positions, seq_time_diffs, targets) in enumerate(tk0):
             # Move tensors to the target device
@@ -152,10 +204,15 @@ class SeqTrainer(object):
             self.optimizer.step()
 
             total_loss += loss.item()
+            epoch_loss += loss.item()
+            batch_count += 1
             if (i + 1) % log_interval == 0:
                 tk0.set_postfix(loss=total_loss / log_interval)
                 total_loss = 0
 
+        # Return average epoch loss
+        return epoch_loss / batch_count if batch_count > 0 else 0
+
     def evaluate(self, data_loader):
         """Evaluate the model on a validation/test data loader.
 
@@ -198,7 +255,7 @@ class SeqTrainer(object):
 
         return avg_loss, accuracy
 
-    def export_onnx(self, output_path, batch_size=2, seq_length=50, vocab_size=None, opset_version=14, dynamic_batch=True, device=None, verbose=False):
+    def export_onnx(self, output_path, batch_size=2, seq_length=50, vocab_size=None, opset_version=14, dynamic_batch=True, device=None, verbose=False, onnx_export_kwargs=None):
         """Export the trained sequence generation model to ONNX format.
 
         This method exports sequence generation models (e.g., HSTU) to ONNX format.
@@ -216,6 +273,7 @@ class SeqTrainer(object):
             device (str, optional): Device for export ('cpu', 'cuda', etc.).
                 If None, defaults to 'cpu' for maximum compatibility.
             verbose (bool): Print export details (default: False).
+            onnx_export_kwargs (dict, optional): Extra kwargs forwarded to ``torch.onnx.export``.
 
         Returns:
             bool: True if export succeeded, False otherwise.
@@ -264,20 +322,38 @@ class SeqTrainer(object):
 
         try:
             with torch.no_grad():
-                torch.onnx.export(
-                    model,
-                    (dummy_seq_tokens,
-                     dummy_seq_time_diffs),
-                    output_path,
-                    input_names=["seq_tokens",
-                                 "seq_time_diffs"],
-                    output_names=["output"],
-                    dynamic_axes=dynamic_axes,
-                    opset_version=opset_version,
-                    do_constant_folding=True,
-                    verbose=verbose,
-                    dynamo=False  # Use legacy exporter for dynamic_axes support
-                )
+                import inspect
+
+                export_kwargs = {
+                    "f": output_path,
+                    "input_names": ["seq_tokens",
+                                    "seq_time_diffs"],
+                    "output_names": ["output"],
+                    "dynamic_axes": dynamic_axes,
+                    "opset_version": opset_version,
+                    "do_constant_folding": True,
+                    "verbose": verbose,
+                }
+
+                if onnx_export_kwargs:
+                    overlap = set(export_kwargs.keys()) & set(onnx_export_kwargs.keys())
+                    overlap.discard("dynamo")
+                    if overlap:
+                        raise ValueError("onnx_export_kwargs contains keys that overlap with explicit args: "
+                                         f"{sorted(overlap)}. Please set them via export_onnx() parameters instead.")
+                    export_kwargs.update(onnx_export_kwargs)
+
+                # Auto-pick exporter:
+                # - dynamic_axes present => prefer legacy exporter (dynamo=False) for dynamic batch/seq
+                # - otherwise prefer dynamo exporter (dynamo=True) on newer torch
+                sig = inspect.signature(torch.onnx.export)
+                if "dynamo" in sig.parameters:
+                    if "dynamo" not in export_kwargs:
+                        export_kwargs["dynamo"] = False if dynamic_axes is not None else True
+                else:
+                    export_kwargs.pop("dynamo", None)
+
+                torch.onnx.export(model, (dummy_seq_tokens, dummy_seq_time_diffs), **export_kwargs)
 
             if verbose:
                 print(f"Successfully exported ONNX model to: {output_path}")

torch_rechub/types.py

@@ -0,0 +1,5 @@
+import os
+import typing as ty
+
+#: Type for path to a file.
+FilePath = ty.Union[str, os.PathLike]