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

torch_rechub/models/ranking/edcn.py

@@ -0,0 +1,117 @@
+"""
+Date: create on 09/13/2022
+References:
+    paper: (KDD'21) EDCN: Enhancing Explicit and Implicit Feature Interactions via Information Sharing for Parallel Deep CTR Models
+    url: https://dlp-kdd.github.io/assets/pdf/DLP-KDD_2021_paper_12.pdf
+Authors: lailai, lailai_zxy@tju.edu.cn
+"""
+
+import torch
+from torch import nn
+from ...basic.layers import LR, MLP, CrossLayer,  EmbeddingLayer
+
+
+class EDCN(torch.nn.Module):
+    """Deep & Cross Network with a mixture of low-rank architecture
+
+    Args:
+        features (list[Feature Class]): training by the whole module.
+        n_cross_layers (int) : the number of layers of feature intersection layers
+        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
+        bridge_type (str): the type interaction function, in ["hadamard_product", "pointwise_addition", "concatenation", "attention_pooling"]
+        use_regulation_module (bool): True, whether to use regulation module
+        temperature (int): the temperature coefficient to control distribution
+    """
+
+    def __init__(self, features, n_cross_layers, mlp_params, bridge_type="hadamard_product", use_regulation_module=True,
+                 temperature=1):
+        super().__init__()
+        self.features = features
+        self.n_cross_layers = n_cross_layers
+        self.num_fields = len(features)
+        self.dims = sum([fea.embed_dim for fea in features])
+        self.fea_dims = [fea.embed_dim for fea in features]
+        self.embedding = EmbeddingLayer(features)
+        self.cross_layers = nn.ModuleList([CrossLayer(self.dims) for _ in range(n_cross_layers)])
+        self.bridge_modules = nn.ModuleList([BridgeModule(self.dims, bridge_type) for _ in range(n_cross_layers)])
+        self.regulation_modules = nn.ModuleList([RegulationModule(self.num_fields,
+                                                   self.fea_dims,
+                                                   tau=temperature,
+                                                   use_regulation=use_regulation_module) for _ in range(n_cross_layers)])
+        mlp_params["dims"] = [self.dims, self.dims]
+        self.mlps = nn.ModuleList([MLP(self.dims, output_layer=False, **mlp_params) for _ in range(n_cross_layers)])
+        self.linear = LR(self.dims * 3)
+
+    def forward(self, x):
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)
+        cross_i, deep_i = self.regulation_modules[0](embed_x)
+        cross_0 = cross_i
+        for i in range(self.n_cross_layers):
+            if i>0:
+                cross_i, deep_i = self.regulation_modules[i](bridge_i)
+            cross_i = cross_i + self.cross_layers[i](cross_0, cross_i)
+            deep_i = self.mlps[i](deep_i)
+            bridge_i = self.bridge_modules[i](cross_i, deep_i)
+        x_stack = torch.cat([cross_i, deep_i, bridge_i], dim=1)
+        y = self.linear(x_stack)
+        return torch.sigmoid(y.squeeze(1))
+
+
+class BridgeModule(torch.nn.Module):
+    def __init__(self, input_dim, bridge_type):
+        super(BridgeModule, self).__init__()
+        assert bridge_type in ["hadamard_product", "pointwise_addition", "concatenation",
+                               "attention_pooling"], 'bridge_type= is not supported'.format(bridge_type)
+        self.bridge_type = bridge_type
+        if bridge_type=="concatenation":
+            self.concat_pooling = nn.Sequential(nn.Linear(input_dim * 2, input_dim),
+                                                nn.ReLU())
+        elif bridge_type=="attention_pooling":
+            self.attention_x = nn.Sequential(nn.Linear(input_dim, input_dim),
+                                             nn.ReLU(),
+                                             nn.Linear(input_dim, input_dim,bias=False),
+                                             nn.Softmax(dim=-1))
+            self.attention_h = nn.Sequential(nn.Linear(input_dim, input_dim),
+                                             nn.ReLU(),
+                                             nn.Linear(input_dim, input_dim,bias=False),
+                                             nn.Softmax(dim=-1))
+    def forward(self, x, h):
+        if self.bridge_type == "hadamard_product":
+            out = x * h
+        elif self.bridge_type == "pointwise_addition":
+            out = x + h
+        elif self.bridge_type == "concatenation":
+            out = self.concat_pooling(torch.cat([x, h], dim=-1))
+        elif self.bridge_type == "attention_pooling":
+            out = self.attention_x(x) * x + self.attention_h(h) * h
+        return out
+
+
+class RegulationModule(torch.nn.Module):
+    def __init__(self, num_fields,
+                 dims,
+                 tau,
+                 use_regulation=True):
+        super(RegulationModule, self).__init__()
+        self.use_regulation = use_regulation
+        if self.use_regulation:
+            self.num_fields = num_fields
+            self.dims = dims
+            self.tau = tau
+            self.g1 = nn.Parameter(torch.ones(num_fields))
+            self.g2 = nn.Parameter(torch.ones(num_fields))
+
+    def forward(self, x):
+        if self.use_regulation:
+            g1 = torch.cat([(self.g1[i]/ self.tau).softmax(dim=-1).unsqueeze(-1).repeat(1, self.dims[i]) for i in range(self.num_fields)], dim=-1)
+            g2 = torch.cat([(self.g2[i] / self.tau).softmax(dim=-1).unsqueeze(-1).repeat(1, self.dims[i]) for i in range(self.num_fields)], dim=-1)
+
+            out1, out2 = g1*x, g2*x
+        else:
+            out1, out2 = x, x
+        return out1, out2
+
+
+
+
+

torch_rechub/models/ranking/fibinet.py

@@ -0,0 +1,50 @@
+"""
+Date: create on 10/19/2022
+References:
+    paper: (RecSys '19) FiBiNET: combining feature importance and bilinear feature interaction for click-through rate prediction
+    url: https://dl.acm.org/doi/abs/10.1145/3298689.3347043
+Authors: lailai, lailai_zxy@tju.edu.cn
+"""
+import torch
+from torch import nn
+from ...basic.layers import  MLP, EmbeddingLayer, SENETLayer, BiLinearInteractionLayer
+from ...basic.features import SparseFeature
+class FiBiNet(torch.nn.Module):
+    """
+        Args:
+        features (list[Feature Class]): training by the whole module.
+        reduction_ratio (int) : Hidden layer reduction factor of SENET layer
+        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
+        bilinear_type (str): the type bilinear interaction function, in ["field_all", "field_each", "field_interaction"], field_all means that all features share a W, field_each means that a feature field corresponds to a W_i, field_interaction means that a feature field intersection corresponds to a W_ij
+    """
+    def __init__(self, features, mlp_params, reduction_ratio=3, bilinear_type="field_interaction", **kwargs):
+        super(FiBiNet, self).__init__()
+        self.features = features
+        self.embedding = EmbeddingLayer(features)
+        embedding_dim = max([fea.embed_dim for fea in features])
+        num_fields = len([fea.embed_dim for fea in features if isinstance(fea, SparseFeature) and fea.shared_with == None])
+        self.senet_layer = SENETLayer(num_fields, reduction_ratio)
+        self.bilinear_interaction = BiLinearInteractionLayer(embedding_dim, num_fields,  bilinear_type)
+        self.dims = num_fields * (num_fields - 1) * embedding_dim
+        self.mlp = MLP(self.dims, **mlp_params)
+
+    def forward(self, x):
+        embed_x = self.embedding(x, self.features)
+        embed_senet = self.senet_layer(embed_x)
+        embed_bi1 = self.bilinear_interaction(embed_x)
+        embed_bi2 = self.bilinear_interaction(embed_senet)
+        shallow_part = torch.flatten(torch.cat([embed_bi1, embed_bi2], dim=1), start_dim=1)
+        mlp_out = self.mlp(shallow_part)
+        return torch.sigmoid(mlp_out.squeeze(1))
+
+
+
+
+
+
+
+
+
+
+
+

torch_rechub/models/ranking/widedeep.py

@@ -1,41 +1,41 @@
-"""
-Date: create on 22/04/2022
-References: 
-    paper: (DLRS'2016) Wide & Deep Learning for Recommender Systems 
-    url: https://arxiv.org/abs/1606.07792
-Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
-"""
-
-import torch
-
-from ...basic.layers import LR, MLP, EmbeddingLayer
-
-
-class WideDeep(torch.nn.Module):
-    """Wide & Deep Learning model.
-
-    Args:
-        wide_features (list): the list of `Feature Class`, training by the wide part module.
-        deep_features (list): the list of `Feature Class`, training by the deep part module.
-        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
-    """
-
-    def __init__(self, wide_features, deep_features, mlp_params):
-        super(WideDeep, self).__init__()
-        self.wide_features = wide_features
-        self.deep_features = deep_features
-        self.wide_dims = sum([fea.embed_dim for fea in wide_features])
-        self.deep_dims = sum([fea.embed_dim for fea in deep_features])
-        self.linear = LR(self.wide_dims)
-        self.embedding = EmbeddingLayer(wide_features + deep_features)
-        self.mlp = MLP(self.deep_dims, **mlp_params)
-
-    def forward(self, x):
-        input_wide = self.embedding(x, self.wide_features, squeeze_dim=True)  #[batch_size, wide_dims]
-        input_deep = self.embedding(x, self.deep_features, squeeze_dim=True)  #[batch_size, deep_dims]
-
-        y_wide = self.linear(input_wide)  #[batch_size, 1]
-        y_deep = self.mlp(input_deep)  #[batch_size, 1]
-        y = y_wide + y_deep
-        y = torch.sigmoid(y.squeeze(1))
-        return y
+"""
+Date: create on 22/04/2022
+References: 
+    paper: (DLRS'2016) Wide & Deep Learning for Recommender Systems 
+    url: https://arxiv.org/abs/1606.07792
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+
+from ...basic.layers import LR, MLP, EmbeddingLayer
+
+
+class WideDeep(torch.nn.Module):
+    """Wide & Deep Learning model.
+
+    Args:
+        wide_features (list): the list of `Feature Class`, training by the wide part module.
+        deep_features (list): the list of `Feature Class`, training by the deep part module.
+        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
+    """
+
+    def __init__(self, wide_features, deep_features, mlp_params):
+        super(WideDeep, self).__init__()
+        self.wide_features = wide_features
+        self.deep_features = deep_features
+        self.wide_dims = sum([fea.embed_dim for fea in wide_features])
+        self.deep_dims = sum([fea.embed_dim for fea in deep_features])
+        self.linear = LR(self.wide_dims)
+        self.embedding = EmbeddingLayer(wide_features + deep_features)
+        self.mlp = MLP(self.deep_dims, **mlp_params)
+
+    def forward(self, x):
+        input_wide = self.embedding(x, self.wide_features, squeeze_dim=True)  #[batch_size, wide_dims]
+        input_deep = self.embedding(x, self.deep_features, squeeze_dim=True)  #[batch_size, deep_dims]
+
+        y_wide = self.linear(input_wide)  #[batch_size, 1]
+        y_deep = self.mlp(input_deep)  #[batch_size, 1]
+        y = y_wide + y_deep
+        y = torch.sigmoid(y.squeeze(1))
+        return y

torch_rechub/trainers/__init__.py

@@ -1,2 +1,3 @@
-from .trainer import CTRTrainer
+from .ctr_trainer import CTRTrainer
+from .match_trainer import MatchTrainer
 from .mtl_trainer import MTLTrainer

torch_rechub/trainers/{trainer.py → ctr_trainer.py}

@@ -1,111 +1,128 @@
-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.
-        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={
-            "lr": 1e-3,
-            "weight_decay": 1e-5
-        },
-        scheduler_fn=None,
-        scheduler_params=None,
-        n_epoch=10,
-        earlystop_patience=10,
-        device="cpu",
-        gpus=[],
-        model_path="./",
-    ):
-        self.model = model  #for uniform weights save method in one gpu or multi gpu
-        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.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.device = torch.device(device)  #torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        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, 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)
-            y_pred = self.model(x_dict)
-            loss = self.criterion(y_pred, y.float())
-            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):
-        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
-            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)
-                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, 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 in enumerate(tk0):
-                x_dict = {k: v.to(self.device) for k, v in x_dict.items()}
-                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
+
+
+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