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

torch_rechub/trainers/ctr_trainer.py

@@ -1,128 +1,191 @@
-import os
-import torch
-import tqdm
-from sklearn.metrics import roc_auc_score
-from ..basic.callback import EarlyStopper
-
-
-class CTRTrainer(object):
-    """A general trainer for single task learning.
-
-    Args:
-        model (nn.Module): any multi task learning model.
-        optimizer_fn (torch.optim): optimizer function of pytorch (default = `torch.optim.Adam`).
-        optimizer_params (dict): parameters of optimizer_fn.
-        scheduler_fn (torch.optim.lr_scheduler) : torch scheduling class, eg. `torch.optim.lr_scheduler.StepLR`.
-        scheduler_params (dict): parameters of optimizer scheduler_fn.
-        n_epoch (int): epoch number of training.
-        earlystop_patience (int): how long to wait after last time validation auc improved (default=10).
-        device (str): `"cpu"` or `"cuda:0"`
-        gpus (list): id of multi gpu (default=[]). If the length >=1, then the model will wrapped by nn.DataParallel.
-        loss_mode (int, optional): the training mode, `{0:point-wise, 1:pair-wise, 2:list-wise}`. Defaults to 0.
-        model_path (str): the path you want to save the model (default="./"). Note only save the best weight in the validation data.
-    """
-
-    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,
-        loss_mode=True,
-        model_path="./",
-    ):
-        self.model = model  # for uniform weights save method in one gpu or multi gpu
-        if gpus is None:
-            gpus = []
-        self.gpus = gpus
-        if len(gpus) > 1:
-            print('parallel running on these gpus:', gpus)
-            self.model = torch.nn.DataParallel(self.model, device_ids=gpus)
-        self.device = torch.device(device)  #torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device)
-        if optimizer_params is None:
-            optimizer_params = {"lr": 1e-3, "weight_decay": 1e-5}
-        self.optimizer = optimizer_fn(self.model.parameters(), **optimizer_params)  #default optimizer
-        self.scheduler = None
-        if scheduler_fn is not None:
-            self.scheduler = scheduler_fn(self.optimizer, **scheduler_params)
-        self.loss_mode = loss_mode
-        self.criterion = torch.nn.BCELoss()  #default loss cross_entropy
-        self.evaluate_fn = roc_auc_score  #default evaluate function
-        self.n_epoch = n_epoch
-        self.early_stopper = EarlyStopper(patience=earlystop_patience)
-        self.model_path = model_path
-
-    def train_one_epoch(self, data_loader, log_interval=10):
-        self.model.train()
-        total_loss = 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
-            y = y.to(self.device).float()
-            if self.loss_mode:
-                y_pred = self.model(x_dict)
-                loss = self.criterion(y_pred, y)
-            else:
-                y_pred, other_loss = self.model(x_dict)
-                loss = self.criterion(y_pred, y) + other_loss
-            self.model.zero_grad()
-            loss.backward()
-            self.optimizer.step()
-            total_loss += loss.item()
-            if (i + 1) % log_interval == 0:
-                tk0.set_postfix(loss=total_loss / log_interval)
-                total_loss = 0
-
-    def fit(self, train_dataloader, val_dataloader=None):
-        for epoch_i in range(self.n_epoch):
-            print('epoch:', epoch_i)
-            self.train_one_epoch(train_dataloader)
-            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
-            if val_dataloader:
-                auc = self.evaluate(self.model, val_dataloader)
-                print('epoch:', epoch_i, 'validation: auc:', auc)
-                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
-
-    def evaluate(self, model, data_loader):
-        model.eval()
-        targets, predicts = list(), list()
-        with torch.no_grad():
-            tk0 = tqdm.tqdm(data_loader, desc="validation", 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()}
-                y = y.to(self.device).float().view(-1, 1)  # 确保y是float类型且维度为[batch_size, 1]
-                if self.loss_mode:
-                    y_pred = model(x_dict)
-                else:
-                    y_pred, _ = model(x_dict)
-                targets.extend(y.tolist())
-                predicts.extend(y_pred.tolist())
-        return self.evaluate_fn(targets, predicts)
-
-    def predict(self, model, data_loader):
-        model.eval()
-        predicts = list()
-        with torch.no_grad():
-            tk0 = tqdm.tqdm(data_loader, desc="predict", 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()}
-                y = y.to(self.device)
-                if self.loss_mode:
-                    y_pred = model(x_dict)
-                else:
-                    y_pred, _ = model(x_dict)
-                predicts.extend(y_pred.tolist())
-        return predicts
+import os
+
+import torch
+import tqdm
+from sklearn.metrics import roc_auc_score
+
+from ..basic.callback import EarlyStopper
+from ..basic.loss_func import RegularizationLoss
+
+
+class CTRTrainer(object):
+    """A general trainer for single task learning.
+
+    Args:
+        model (nn.Module): any multi task learning model.
+        optimizer_fn (torch.optim): optimizer function of pytorch (default = `torch.optim.Adam`).
+        optimizer_params (dict): parameters of optimizer_fn.
+        scheduler_fn (torch.optim.lr_scheduler) : torch scheduling class, eg. `torch.optim.lr_scheduler.StepLR`.
+        scheduler_params (dict): parameters of optimizer scheduler_fn.
+        n_epoch (int): epoch number of training.
+        earlystop_patience (int): how long to wait after last time validation auc improved (default=10).
+        device (str): `"cpu"` or `"cuda:0"`
+        gpus (list): id of multi gpu (default=[]). If the length >=1, then the model will wrapped by nn.DataParallel.
+        loss_mode (int, optional): the training mode, `{0:point-wise, 1:pair-wise, 2:list-wise}`. Defaults to 0.
+        model_path (str): the path you want to save the model (default="./"). Note only save the best weight in the validation data.
+        embedding_l1 (float): L1 regularization coefficient for embedding parameters (default=0.0).
+        embedding_l2 (float): L2 regularization coefficient for embedding parameters (default=0.0).
+        dense_l1 (float): L1 regularization coefficient for dense parameters (default=0.0).
+        dense_l2 (float): L2 regularization coefficient for dense parameters (default=0.0).
+    """
+
+    def __init__(
+        self,
+        model,
+        optimizer_fn=torch.optim.Adam,
+        optimizer_params=None,
+        regularization_params=None,
+        scheduler_fn=None,
+        scheduler_params=None,
+        n_epoch=10,
+        earlystop_patience=10,
+        device="cpu",
+        gpus=None,
+        loss_mode=True,
+        model_path="./",
+    ):
+        self.model = model  # for uniform weights save method in one gpu or multi gpu
+        if gpus is None:
+            gpus = []
+        self.gpus = gpus
+        if len(gpus) > 1:
+            print('parallel running on these gpus:', gpus)
+            self.model = torch.nn.DataParallel(self.model, device_ids=gpus)
+        # torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device(device)
+        self.model.to(self.device)
+        if optimizer_params is None:
+            optimizer_params = {"lr": 1e-3, "weight_decay": 1e-5}
+        self.optimizer = optimizer_fn(self.model.parameters(), **optimizer_params)  # default optimizer
+        if regularization_params is None:
+            regularization_params = {"embedding_l1": 0.0, "embedding_l2": 0.0, "dense_l1": 0.0, "dense_l2": 0.0}
+        self.scheduler = None
+        if scheduler_fn is not None:
+            self.scheduler = scheduler_fn(self.optimizer, **scheduler_params)
+        self.loss_mode = loss_mode
+        self.criterion = torch.nn.BCELoss()  # default loss cross_entropy
+        self.evaluate_fn = roc_auc_score  # default evaluate function
+        self.n_epoch = n_epoch
+        self.early_stopper = EarlyStopper(patience=earlystop_patience)
+        self.model_path = model_path
+        # Initialize regularization loss
+        self.reg_loss_fn = RegularizationLoss(**regularization_params)
+
+    def train_one_epoch(self, data_loader, log_interval=10):
+        self.model.train()
+        total_loss = 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
+            y = y.to(self.device).float()
+            if self.loss_mode:
+                y_pred = self.model(x_dict)
+                loss = self.criterion(y_pred, y)
+            else:
+                y_pred, other_loss = self.model(x_dict)
+                loss = self.criterion(y_pred, y) + other_loss
+
+            # Add regularization loss
+            reg_loss = self.reg_loss_fn(self.model)
+            loss = loss + reg_loss
+
+            self.model.zero_grad()
+            loss.backward()
+            self.optimizer.step()
+            total_loss += loss.item()
+            if (i + 1) % log_interval == 0:
+                tk0.set_postfix(loss=total_loss / log_interval)
+                total_loss = 0
+
+    def fit(self, train_dataloader, val_dataloader=None):
+        for epoch_i in range(self.n_epoch):
+            print('epoch:', epoch_i)
+            self.train_one_epoch(train_dataloader)
+            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
+            if val_dataloader:
+                auc = self.evaluate(self.model, val_dataloader)
+                print('epoch:', epoch_i, 'validation: auc:', auc)
+                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
+
+    def evaluate(self, model, data_loader):
+        model.eval()
+        targets, predicts = list(), list()
+        with torch.no_grad():
+            tk0 = tqdm.tqdm(data_loader, desc="validation", 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()}
+                # 确保y是float类型且维度为[batch_size, 1]
+                y = y.to(self.device).float().view(-1, 1)
+                if self.loss_mode:
+                    y_pred = model(x_dict)
+                else:
+                    y_pred, _ = model(x_dict)
+                targets.extend(y.tolist())
+                predicts.extend(y_pred.tolist())
+        return self.evaluate_fn(targets, predicts)
+
+    def predict(self, model, data_loader):
+        model.eval()
+        predicts = list()
+        with torch.no_grad():
+            tk0 = tqdm.tqdm(data_loader, desc="predict", 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()}
+                y = y.to(self.device)
+                if self.loss_mode:
+                    y_pred = model(x_dict)
+                else:
+                    y_pred, _ = model(x_dict)
+                predicts.extend(y_pred.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):
+        """Export the trained model to ONNX format.
+
+        This method exports the ranking model (e.g., DeepFM, WideDeep, DCN) to ONNX format
+        for deployment. The export is non-invasive and does not modify the model code.
+
+        Args:
+            output_path (str): Path to save the ONNX model file.
+            dummy_input (dict, optional): Example input dict {feature_name: tensor}.
+                If not provided, dummy inputs will be generated automatically.
+            batch_size (int): Batch size for auto-generated dummy input (default: 2).
+            seq_length (int): Sequence length for SequenceFeature (default: 10).
+            opset_version (int): ONNX opset version (default: 14).
+            dynamic_batch (bool): Enable dynamic batch size (default: True).
+            device (str, optional): Device for export ('cpu', 'cuda', etc.).
+                If None, defaults to 'cpu' for maximum compatibility.
+            verbose (bool): Print export details (default: False).
+
+        Returns:
+            bool: True if export succeeded, False otherwise.
+
+        Example:
+            >>> trainer = CTRTrainer(model, ...)
+            >>> trainer.fit(train_dl, val_dl)
+            >>> trainer.export_onnx("deepfm.onnx")
+
+            >>> # With custom dummy input
+            >>> dummy = {"user_id": torch.tensor([1, 2]), "item_id": torch.tensor([10, 20])}
+            >>> trainer.export_onnx("model.onnx", dummy_input=dummy)
+
+            >>> # Export on specific device
+            >>> trainer.export_onnx("model.onnx", device="cpu")
+        """
+        from ..utils.onnx_export import ONNXExporter
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        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)