torch-rechub 0.0.4__py3-none-any.whl → 0.0.6__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()
@@ -237,3 +275,100 @@ class MatchTrainer(object):
             # Restore original mode
             if hasattr(model, 'mode'):
                 model.mode = original_mode
+
+    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.
+
+        This method generates a visual representation of the model architecture,
+        showing layer connections, tensor shapes, and nested module structures.
+        It automatically extracts feature information from the model.
+
+        Parameters
+        ----------
+        input_data : dict, optional
+            Example input dict {feature_name: tensor}.
+            If not provided, dummy inputs will be generated automatically.
+        batch_size : int, default=2
+            Batch size for auto-generated dummy input.
+        seq_length : int, default=10
+            Sequence length for SequenceFeature.
+        depth : int, default=3
+            Visualization depth, higher values show more detail.
+            Set to -1 to show all layers.
+        show_shapes : bool, default=True
+            Whether to display tensor shapes.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        dpi : int, default=300
+            Resolution in dots per inch for output image.
+            Higher values produce sharper images suitable for papers.
+        **kwargs : dict
+            Additional arguments passed to torchview.draw_graph().
+
+        Returns
+        -------
+        ComputationGraph
+            A torchview ComputationGraph object.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+
+        Notes
+        -----
+        Default Display Behavior:
+            When `save_path` is None (default):
+            - In Jupyter/IPython: automatically displays the graph inline
+            - In Python script: opens the graph with system default viewer
+
+        Examples
+        --------
+        >>> trainer = MatchTrainer(model, ...)
+        >>> trainer.fit(train_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        from ..utils.visualization import TORCHVIEW_AVAILABLE, visualize_model
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "Install with: pip install torch-rechub[visualization]\n"
+                "Also ensure graphviz is installed on your system:\n"
+                "  - Ubuntu/Debian: sudo apt-get install graphviz\n"
+                "  - macOS: brew install graphviz\n"
+                "  - Windows: choco install graphviz"
+            )
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        return visualize_model(
+            model,
+            input_data=input_data,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            depth=depth,
+            show_shapes=show_shapes,
+            expand_nested=expand_nested,
+            save_path=save_path,
+            graph_name=graph_name,
+            device=viz_device,
+            dpi=dpi,
+            **kwargs
+        )

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()
@@ -257,3 +305,100 @@ class MTLTrainer(object):
 
         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)
+
+    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.
+
+        This method generates a visual representation of the model architecture,
+        showing layer connections, tensor shapes, and nested module structures.
+        It automatically extracts feature information from the model.
+
+        Parameters
+        ----------
+        input_data : dict, optional
+            Example input dict {feature_name: tensor}.
+            If not provided, dummy inputs will be generated automatically.
+        batch_size : int, default=2
+            Batch size for auto-generated dummy input.
+        seq_length : int, default=10
+            Sequence length for SequenceFeature.
+        depth : int, default=3
+            Visualization depth, higher values show more detail.
+            Set to -1 to show all layers.
+        show_shapes : bool, default=True
+            Whether to display tensor shapes.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        dpi : int, default=300
+            Resolution in dots per inch for output image.
+            Higher values produce sharper images suitable for papers.
+        **kwargs : dict
+            Additional arguments passed to torchview.draw_graph().
+
+        Returns
+        -------
+        ComputationGraph
+            A torchview ComputationGraph object.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+
+        Notes
+        -----
+        Default Display Behavior:
+            When `save_path` is None (default):
+            - In Jupyter/IPython: automatically displays the graph inline
+            - In Python script: opens the graph with system default viewer
+
+        Examples
+        --------
+        >>> trainer = MTLTrainer(model, task_types=["classification", "classification"])
+        >>> trainer.fit(train_dl, val_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        from ..utils.visualization import TORCHVIEW_AVAILABLE, visualize_model
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "Install with: pip install torch-rechub[visualization]\n"
+                "Also ensure graphviz is installed on your system:\n"
+                "  - Ubuntu/Debian: sudo apt-get install graphviz\n"
+                "  - macOS: brew install graphviz\n"
+                "  - Windows: choco install graphviz"
+            )
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        return visualize_model(
+            model,
+            input_data=input_data,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            depth=depth,
+            show_shapes=show_shapes,
+            expand_nested=expand_nested,
+            save_path=save_path,
+            graph_name=graph_name,
+            device=viz_device,
+            dpi=dpi,
+            **kwargs
+        )

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.
 
@@ -291,3 +348,137 @@ class SeqTrainer(object):
         except Exception as e:
             warnings.warn(f"ONNX export failed: {str(e)}")
             raise RuntimeError(f"Failed to export ONNX model: {str(e)}") from e
+
+    def visualization(self, seq_length=50, vocab_size=None, batch_size=2, 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.
+
+        This method generates a visual representation of the sequence model
+        architecture, showing layer connections, tensor shapes, and nested
+        module structures.
+
+        Parameters
+        ----------
+        seq_length : int, default=50
+            Sequence length for dummy input.
+        vocab_size : int, optional
+            Vocabulary size for generating dummy tokens.
+            If None, will try to get from model.vocab_size or model.item_num.
+        batch_size : int, default=2
+            Batch size for dummy input.
+        depth : int, default=3
+            Visualization depth, higher values show more detail.
+            Set to -1 to show all layers.
+        show_shapes : bool, default=True
+            Whether to display tensor shapes.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        dpi : int, default=300
+            Resolution in dots per inch for output image.
+            Higher values produce sharper images suitable for papers.
+        **kwargs : dict
+            Additional arguments passed to torchview.draw_graph().
+
+        Returns
+        -------
+        ComputationGraph
+            A torchview ComputationGraph object.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+        ValueError
+            If vocab_size is not provided and cannot be inferred from model.
+
+        Notes
+        -----
+        Default Display Behavior:
+            When `save_path` is None (default):
+            - In Jupyter/IPython: automatically displays the graph inline
+            - In Python script: opens the graph with system default viewer
+
+        Examples
+        --------
+        >>> trainer = SeqTrainer(hstu_model, ...)
+        >>> trainer.fit(train_dl, val_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4, vocab_size=10000)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        try:
+            from torchview import draw_graph
+            TORCHVIEW_AVAILABLE = True
+        except ImportError:
+            TORCHVIEW_AVAILABLE = False
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "Install with: pip install torch-rechub[visualization]\n"
+                "Also ensure graphviz is installed on your system:\n"
+                "  - Ubuntu/Debian: sudo apt-get install graphviz\n"
+                "  - macOS: brew install graphviz\n"
+                "  - Windows: choco install graphviz"
+            )
+
+        from ..utils.visualization import _is_jupyter_environment, display_graph
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        # Get vocab_size from model if not provided
+        if vocab_size is None:
+            if hasattr(model, 'vocab_size'):
+                vocab_size = model.vocab_size
+            elif hasattr(model, 'item_num'):
+                vocab_size = model.item_num
+            else:
+                raise ValueError("vocab_size must be provided or model must have "
+                                 "'vocab_size' or 'item_num' attribute")
+
+        # Generate dummy inputs for sequence model
+        dummy_seq_tokens = torch.randint(0, vocab_size, (batch_size, seq_length), device=viz_device)
+        dummy_seq_time_diffs = torch.zeros(batch_size, seq_length, dtype=torch.float32, device=viz_device)
+
+        # Move model to device
+        model = model.to(viz_device)
+        model.eval()
+
+        # Call torchview.draw_graph
+        graph = draw_graph(model, input_data=(dummy_seq_tokens, dummy_seq_time_diffs), graph_name=graph_name, depth=depth, device=viz_device, expand_nested=expand_nested, show_shapes=show_shapes, save_graph=False, **kwargs)
+
+        # Set DPI for high-quality output
+        graph.visual_graph.graph_attr['dpi'] = str(dpi)
+
+        # Handle save_path: manually save with DPI applied
+        if save_path:
+            import os
+            directory = os.path.dirname(save_path) or "."
+            filename = os.path.splitext(os.path.basename(save_path))[0]
+            ext = os.path.splitext(save_path)[1].lstrip('.')
+            output_format = ext if ext else 'pdf'
+            if directory != "." and not os.path.exists(directory):
+                os.makedirs(directory, exist_ok=True)
+            graph.visual_graph.render(filename=filename, directory=directory, format=output_format, cleanup=True)
+
+        # Handle default display behavior when save_path is None
+        if save_path is None:
+            if _is_jupyter_environment():
+                display_graph(graph)
+            else:
+                graph.visual_graph.view(cleanup=True)
+
+        return graph