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

torch_rechub/trainers/match_trainer.py

@@ -0,0 +1,170 @@
+import os
+import torch
+import tqdm
+from sklearn.metrics import roc_auc_score
+from ..basic.callback import EarlyStopper
+from ..basic.loss_func import BPRLoss
+
+
+class MatchTrainer(object):
+    """A general trainer for Matching/Retrieval
+
+    Args:
+        model (nn.Module): any matching model.
+        mode (int, optional): the training mode, `{0:point-wise, 1:pair-wise, 2:list-wise}`. Defaults to 0.
+        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.
+        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,
+        mode=0,
+        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="./",
+    ):
+        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.mode = mode
+        if mode == 0:  #point-wise loss, binary cross_entropy
+            self.criterion = torch.nn.BCELoss()  #default loss binary cross_entropy
+        elif mode == 1:  #pair-wise loss
+            self.criterion = BPRLoss()
+        elif mode == 2:  #list-wise loss, softmax
+            self.criterion = torch.nn.CrossEntropyLoss()
+        else:
+            raise ValueError("mode only contain value in %s, but got %s" % ([0, 1, 2], mode))
+        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.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)
+            if self.mode == 0:
+                y = y.float()  #torch._C._nn.binary_cross_entropy expected Float
+            else:
+                y = y.long()  #
+            if self.mode == 1:  #pair_wise
+                pos_score, neg_score = self.model(x_dict)
+                loss = self.criterion(pos_score, neg_score)
+            else:
+                y_pred = self.model(x_dict)
+                loss = self.criterion(y_pred, y)
+            # used for debug
+            # if i == 0:
+            #     print()
+            #     if self.mode == 0:
+            #         print('pred: ', [f'{float(each):5.2g}' for each in y_pred.detach().cpu().tolist()])
+            #         print('truth:', [f'{float(each):5.2g}' for each in y.detach().cpu().tolist()])
+            #     elif self.mode == 2:
+            #         pred = y_pred.detach().cpu().mean(0)
+            #         pred = torch.softmax(pred, dim=0).tolist()
+            #         print('pred: ', [f'{float(each):4.2g}' for each in pred])
+            #     elif self.mode == 1:
+            #         print('pos:', [f'{float(each):5.2g}' for each in pos_score.detach().cpu().tolist()])
+            #         print('neg: ', [f'{float(each):5.2g}' for each in neg_score.detach().cpu().tolist()])
+
+            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)
+                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)
+                y_pred = model(x_dict)
+                predicts.extend(y_pred.tolist())
+        return predicts
+
+    def inference_embedding(self, model, mode, data_loader, model_path):
+        #inference
+        assert mode in ["user", "item"], "Invalid mode={}.".format(mode)
+        model.mode = mode
+        model.load_state_dict(torch.load(os.path.join(model_path, "model.pth")))
+        model = model.to(self.device)
+        model.eval()
+        predicts = []
+        with torch.no_grad():
+            tk0 = tqdm.tqdm(data_loader, desc="%s inference" % (mode), smoothing=0, mininterval=1.0)
+            for i, x_dict in enumerate(tk0):
+                x_dict = {k: v.to(self.device) for k, v in x_dict.items()}
+                y_pred = model(x_dict)
+                predicts.append(y_pred.data)
+        return torch.cat(predicts, dim=0)

torch_rechub/trainers/mtl_trainer.py

@@ -1,145 +1,207 @@
-import os
-import tqdm
-import numpy as np
-import torch
-import torch.nn as nn
-from ..basic.callback import EarlyStopper
-from ..basic.utils import get_loss_func, get_metric_func
-from ..models.multi_task import ESMM
-
-
-class MTLTrainer(object):
-    """A trainer for multi task learning.
-
-    Args:
-        model (nn.Module): any multi task learning model.
-        task_types (list): types of tasks, only support ["classfication", "regression"].
-        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.
-        adaptive_params (dict): parameters of adaptive loss weight method. Now only support `{"method" : "uwl"}`. 
-        n_epoch (int): epoch number of training.
-        earlystop_taskid (int): task id of earlystop metrics relies between multi task (default = 0).
-        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.
-        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,
-        task_types,
-        optimizer_fn=torch.optim.Adam,
-        optimizer_params={
-            "lr": 1e-3,
-            "weight_decay": 1e-5
-        },
-        scheduler_fn=None,
-        scheduler_params=None,
-        adaptive_params=None,
-        n_epoch=10,
-        earlystop_taskid=0,
-        earlystop_patience=10,
-        device="cpu",
-        gpus=[],
-        model_path="./",
-    ):
-        self.model = model
-        self.task_types = task_types
-        self.n_task = len(task_types)
-        self.loss_weight = None
-        self.adaptive_method = None
-        if adaptive_params is not None:
-            if adaptive_params["method"] == "uwl":
-                self.adaptive_method = "uwl"
-                self.loss_weight = nn.ParameterList(nn.Parameter(torch.zeros(1)) for _ in range(self.n_task))
-                self.model.add_module("loss weight", self.loss_weight)
-        self.optimizer = optimizer_fn(self.model.parameters(), **optimizer_params)  #default Adam optimizer
-        self.scheduler = None
-        if scheduler_fn is not None:
-            self.scheduler = scheduler_fn(self.optimizer, **scheduler_params)
-        self.loss_fns = [get_loss_func(task_type) for task_type in task_types]
-        self.evaluate_fns = [get_metric_func(task_type) for task_type in task_types]
-        self.n_epoch = n_epoch
-        self.earlystop_taskid = earlystop_taskid
-        self.early_stopper = EarlyStopper(patience=earlystop_patience)
-        self.device = torch.device(device)
-
-        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.model_path = model_path
-
-    def train_one_epoch(self, data_loader):
-        self.model.train()
-        total_loss = np.zeros(self.n_task)
-        tk0 = tqdm.tqdm(data_loader, desc="train", smoothing=0, mininterval=1.0)
-        for iter_i, (x_dict, ys) in enumerate(tk0):
-            x_dict = {k: v.to(self.device) for k, v in x_dict.items()}  #tensor to GPU
-            ys = ys.to(self.device)
-            y_preds = self.model(x_dict)
-            loss_list = [self.loss_fns[i](y_preds[:, i], ys[:, i].float()) for i in range(self.n_task)]
-            if isinstance(self.model, ESMM):
-                loss = sum(loss_list[1:])  #ESSM only compute loss for ctr and ctcvr task
-            else:
-                if self.adaptive_method != None:
-                    if self.adaptive_method == "uwl":
-                        loss = 0
-                        for loss_i, w_i in zip(loss_list, self.loss_weight):
-                            w_i = torch.clamp(w_i, min=0)
-                            loss += 2 * loss_i * torch.exp(-w_i) + w_i
-                else:
-                    loss = sum(loss_list) / self.n_task
-            self.model.zero_grad()
-            loss.backward()
-            self.optimizer.step()
-            total_loss += np.array([l.item() for l in loss_list])
-        log_dict = {"task_%d:" % (i): total_loss[i] / (iter_i + 1) for i in range(self.n_task)}
-        print("train loss: ", log_dict)
-        if self.loss_weight:
-            print("loss weight: ", [w.item() for w in self.loss_weight])
-
-    def fit(self, train_dataloader, val_dataloader):
-        self.model.to(self.device)
-        for epoch_i in range(self.n_epoch):
-            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
-            scores = self.evaluate(self.model, val_dataloader)
-            print('epoch:', epoch_i, 'validation scores: ', scores)
-            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)
-                torch.save(self.early_stopper.best_weights, os.path.join(self.model_path, "model.pth"))  #save best auc model
-                break
-
-    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, ys) in enumerate(tk0):
-                x_dict = {k: v.to(self.device) for k, v in x_dict.items()}  #tensor to GPU
-                ys = ys.to(self.device)
-                y_preds = self.model(x_dict)
-                targets.extend(ys.tolist())
-                predicts.extend(y_preds.tolist())
-        targets, predicts = np.array(targets), np.array(predicts)
-        scores = [self.evaluate_fns[i](targets[:, i], predicts[:, i]) for i in range(self.n_task)]
-        return scores
-
-    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 in enumerate(tk0):
-                x_dict = {k: v.to(self.device) for k, v in x_dict.items()}
-                y_preds = model(x_dict)
-                predicts.extend(y_preds.tolist())
+import os
+import tqdm
+import numpy as np
+import torch
+import torch.nn as nn
+from ..basic.callback import EarlyStopper
+from ..utils.data import get_loss_func, get_metric_func
+from ..models.multi_task import ESMM
+from ..utils.mtl import shared_task_layers, gradnorm, MetaBalance
+
+
+class MTLTrainer(object):
+    """A trainer for multi task learning.
+
+    Args:
+        model (nn.Module): any multi task learning model.
+        task_types (list): types of tasks, only support ["classfication", "regression"].
+        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.
+        adaptive_params (dict): parameters of adaptive loss weight method. Now only support `{"method" : "uwl"}`. 
+        n_epoch (int): epoch number of training.
+        earlystop_taskid (int): task id of earlystop metrics relies between multi task (default = 0).
+        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.
+        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,
+        task_types,
+        optimizer_fn=torch.optim.Adam,
+        optimizer_params=None,
+        scheduler_fn=None,
+        scheduler_params=None,
+        adaptive_params=None,
+        n_epoch=10,
+        earlystop_taskid=0,
+        earlystop_patience=10,
+        device="cpu",
+        gpus=None,
+        model_path="./",
+    ):
+        self.model = model
+        if gpus is None:
+            gpus = []
+        if optimizer_params is None:
+            optimizer_params = {
+                "lr": 1e-3,
+                "weight_decay": 1e-5
+            }
+        self.task_types = task_types
+        self.n_task = len(task_types)
+        self.loss_weight = None
+        self.adaptive_method = None
+        if adaptive_params is not None:
+            if adaptive_params["method"] == "uwl":
+                self.adaptive_method = "uwl"
+                self.loss_weight = nn.ParameterList(nn.Parameter(torch.zeros(1)) for _ in range(self.n_task))
+                self.model.add_module("loss weight", self.loss_weight)
+            elif adaptive_params["method"] == "metabalance":
+                self.adaptive_method = "metabalance"
+                share_layers, task_layers = shared_task_layers(self.model)
+                self.meta_optimizer = MetaBalance(share_layers)
+                self.share_optimizer = optimizer_fn(share_layers, **optimizer_params)
+                self.task_optimizer = optimizer_fn(task_layers, **optimizer_params)
+            elif adaptive_params["method"] == "gradnorm":
+                self.adaptive_method = "gradnorm"
+                self.alpha = adaptive_params.get("alpha", 0.16)
+                share_layers = shared_task_layers(self.model)[0]
+                #gradnorm calculate the gradients of each loss on the last fully connected shared layer weight(dimension is 2)
+                for i in range(len(share_layers)):
+                    if share_layers[-i].ndim == 2:
+                        self.last_share_layer = share_layers[-i]
+                        break
+                self.initial_task_loss = None
+                self.loss_weight = nn.ParameterList(nn.Parameter(torch.ones(1)) for _ in range(self.n_task))
+                self.model.add_module("loss weight", self.loss_weight)
+        if self.adaptive_method != "metabalance":
+            self.optimizer = optimizer_fn(self.model.parameters(), **optimizer_params)  #default Adam optimizer
+        self.scheduler = None
+        if scheduler_fn is not None:
+            self.scheduler = scheduler_fn(self.optimizer, **scheduler_params)
+        self.loss_fns = [get_loss_func(task_type) for task_type in task_types]
+        self.evaluate_fns = [get_metric_func(task_type) for task_type in task_types]
+        self.n_epoch = n_epoch
+        self.earlystop_taskid = earlystop_taskid
+        self.early_stopper = EarlyStopper(patience=earlystop_patience)
+
+        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)
+        self.model_path = model_path
+
+
+    def train_one_epoch(self, data_loader):
+        self.model.train()
+        total_loss = np.zeros(self.n_task)
+        tk0 = tqdm.tqdm(data_loader, desc="train", smoothing=0, mininterval=1.0)
+        for iter_i, (x_dict, ys) in enumerate(tk0):
+            x_dict = {k: v.to(self.device) for k, v in x_dict.items()}  #tensor to GPU
+            ys = ys.to(self.device)
+            y_preds = self.model(x_dict)
+            loss_list = [self.loss_fns[i](y_preds[:, i], ys[:, i].float()) for i in range(self.n_task)]
+            if isinstance(self.model, ESMM):
+                loss = sum(loss_list[1:])  #ESSM only compute loss for ctr and ctcvr task
+            else:
+                if self.adaptive_method != None:
+                    if self.adaptive_method == "uwl":
+                        loss = 0
+                        for loss_i, w_i in zip(loss_list, self.loss_weight):
+                            w_i = torch.clamp(w_i, min=0)
+                            loss += 2 * loss_i * torch.exp(-w_i) + w_i
+                else:
+                    loss = sum(loss_list) / self.n_task
+            if self.adaptive_method == 'metabalance':
+                self.share_optimizer.zero_grad()
+                self.task_optimizer.zero_grad()
+                self.meta_optimizer.step(loss_list)
+                self.share_optimizer.step()
+                self.task_optimizer.step()
+            elif self.adaptive_method == "gradnorm":
+                self.optimizer.zero_grad()
+                if self.initial_task_loss is None:
+                    self.initial_task_loss = [l.item() for l in loss_list]
+                gradnorm(loss_list, self.loss_weight, self.last_share_layer, self.initial_task_loss, self.alpha)
+                self.optimizer.step()
+                # renormalize
+                loss_weight_sum = sum([w.item() for w in self.loss_weight])
+                normalize_coeff = len(self.loss_weight) / loss_weight_sum
+                for w in self.loss_weight:
+                    w.data = w.data * normalize_coeff
+            else:
+                self.model.zero_grad()
+                loss.backward()
+                self.optimizer.step()
+            total_loss += np.array([l.item() for l in loss_list])
+        log_dict = {"task_%d:" % (i): total_loss[i] / (iter_i + 1) for i in range(self.n_task)}
+        loss_list = [total_loss[i] / (iter_i + 1) for i in range(self.n_task)]
+        print("train loss: ", log_dict)
+        if self.loss_weight:
+            print("loss weight: ", [w.item() for w in self.loss_weight])
+
+        return loss_list
+
+
+    def fit(self, train_dataloader, val_dataloader, mode = 'base', seed = 0):
+        total_log = []
+
+        for epoch_i in range(self.n_epoch):
+            _log_per_epoch = 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
+            scores = self.evaluate(self.model, val_dataloader)
+            print('epoch:', epoch_i, 'validation scores: ', scores)
+
+            for score in scores:
+                _log_per_epoch.append(score)
+
+            total_log.append(_log_per_epoch)
+
+            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)
+                break
+
+        torch.save(self.model.state_dict(), os.path.join(self.model_path, "model_{}_{}.pth".format(mode, seed)))  #save best auc model
+
+        return total_log
+
+
+    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, ys) in enumerate(tk0):
+                x_dict = {k: v.to(self.device) for k, v in x_dict.items()}  #tensor to GPU
+                ys = ys.to(self.device)
+                y_preds = self.model(x_dict)
+                targets.extend(ys.tolist())
+                predicts.extend(y_preds.tolist())
+        targets, predicts = np.array(targets), np.array(predicts)
+        scores = [self.evaluate_fns[i](targets[:, i], predicts[:, i]) for i in range(self.n_task)]
+        return scores
+
+
+    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 in enumerate(tk0):
+                x_dict = {k: v.to(self.device) for k, v in x_dict.items()}
+                y_preds = model(x_dict)
+                predicts.extend(y_preds.tolist())
         return predicts