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

torch_rechub/models/ranking/dien.py

@@ -0,0 +1,191 @@
+"""
+Date: create on 01/05/2024
+References:
+    paper: (AAAI'2019) Deep Interest Evolution Network for Click-Through Rate Prediction
+    url: https://arxiv.org/pdf/1809.03672
+Authors: Tao Fan, thisisevy@foxmail.com
+"""
+
+import torch
+from torch import nn
+from torch.nn import Parameter, init
+
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class AUGRU(nn.Module):
+
+    def __init__(self, embed_dim):
+        super(AUGRU, self).__init__()
+        self.embed_dim = embed_dim
+        # 初始化AUGRU单元
+        self.augru_cell = AUGRU_Cell(self.embed_dim)
+
+    def forward(self, x, item):
+        '''
+        :param x: 输入的序列向量，维度为 [ batch_size, seq_lens, embed_dim ]
+        :param item: 目标物品的向量
+        :return: outs: 所有AUGRU单元输出的隐藏向量[ batch_size, seq_lens, embed_dim ]
+                 h: 最后一个AUGRU单元输出的隐藏向量[ batch_size, embed_dim ]
+        '''
+        outs = []
+        h = None
+        # 开始循环，x.shape[1]是序列的长度
+        for i in range(x.shape[1]):
+            if h is None:
+                # 初始化第一层的输入h
+                h = Parameter(torch.rand(x.shape[0], self.embed_dim).to(x.device))
+            h = self.augru_cell(x[:, i], h, item)
+            outs.append(torch.unsqueeze(h, dim=1))
+        outs = torch.cat(outs, dim=1)
+        return outs, h
+
+
+# AUGRU单元
+class AUGRU_Cell(nn.Module):
+
+    def __init__(self, embed_dim):
+        """
+        :param embed_dim: 输入向量的维度
+        """
+        super(AUGRU_Cell, self).__init__()
+
+        # 初始化更新门的模型参数
+        self.Wu = Parameter(torch.rand(embed_dim, embed_dim))
+        self.Uu = Parameter(torch.rand(embed_dim, embed_dim))
+        self.bu = init.xavier_uniform_(Parameter(torch.empty(1, embed_dim)))
+
+        # 初始化重置门的模型参数
+        self.Wr = init.xavier_uniform_(Parameter(torch.empty(embed_dim, embed_dim)))
+        self.Ur = init.xavier_uniform_(Parameter(torch.empty(embed_dim, embed_dim)))
+        self.br = init.xavier_uniform_(Parameter(torch.empty(1, embed_dim)))
+
+        # 初始化计算h~的模型参数
+        self.Wh = init.xavier_uniform_(Parameter(torch.empty(embed_dim, embed_dim)))
+        self.Uh = init.xavier_uniform_(Parameter(torch.empty(embed_dim, embed_dim)))
+        self.bh = init.xavier_uniform_(Parameter(torch.empty(1, embed_dim)))
+
+        # 初始化注意计算里的模型参数
+        self.Wa = init.xavier_uniform_(Parameter(torch.empty(embed_dim, embed_dim)))
+
+
+# 注意力的计算
+
+    def attention(self, x, item):
+        '''
+        :param x: 输入的序列中第t个向量 [ batch_size, embed_dim ]
+        :param item: 目标物品的向量 [ batch_size, embed_dim ]
+        :return: 注意力权重 [ batch_size, 1 ]
+        '''
+        hW = torch.matmul(x, self.Wa)
+        hWi = torch.sum(hW * item, dim=1)
+        hWi = torch.unsqueeze(hWi, 1)
+        return torch.softmax(hWi, dim=1)
+
+    def forward(self, x, h_1, item):
+        '''
+        :param x:  输入的序列中第t个物品向量 [ batch_size, embed_dim ]
+        :param h_1:  上一个AUGRU单元输出的隐藏向量 [ batch_size, embed_dim ]
+        :param item: 目标物品的向量 [ batch_size, embed_dim ]
+        :return: h 当前层输出的隐藏向量 [ batch_size, embed_dim ]
+        '''
+        # [ batch_size, embed_dim ]
+        u = torch.sigmoid(torch.matmul(x, self.Wu) + torch.matmul(h_1, self.Uu) + self.bu)
+        # [ batch_size, embed_dim ]
+        r = torch.sigmoid(torch.matmul(x, self.Wr) + torch.matmul(h_1, self.Ur) + self.br)
+        # [ batch_size, embed_dim ]
+        h_hat = torch.tanh(torch.matmul(x, self.Wh) + r * torch.matmul(h_1, self.Uh) + self.bh)
+        # [ batch_size, 1 ]
+        a = self.attention(x, item)
+        # [ batch_size, embed_dim ]
+        u_hat = a * u
+        # [ batch_size, embed_dim ]
+        h = (1 - u_hat) * h_1 + u_hat * h_hat
+        # [ batch_size, embed_dim ]
+        return h
+
+
+class DIEN(nn.Module):
+    """Deep Interest Evolution Network
+    Args:
+        features (list): the list of `Feature Class`. training by MLP. It means the user profile features and context features in origin paper, exclude history and target features.
+        history_features (list): the list of `Feature Class`,training by ActivationUnit. It means the user behaviour sequence features, eg.item id sequence, shop id sequence.
+        target_features (list): the list of `Feature Class`, training by ActivationUnit. It means the target feature which will execute target-attention with history feature.
+        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
+        history_labels (list): the list of history_features whether it is clicked history or not. It should be 0 or 1.
+        alpha (float): the weighting of auxiliary loss.
+    """
+
+    def __init__(self, features, history_features, target_features, mlp_params, history_labels, alpha=0.2):
+        super().__init__()
+        self.alpha = alpha  # 计算辅助损失函数时的权重
+        self.features = features
+        self.history_features = history_features
+        self.target_features = target_features
+        self.num_history_features = len(history_features)
+        self.all_dims = sum([fea.embed_dim for fea in features + history_features + target_features])
+        # self.GRU = nn.GRU(batch_first=True)
+        self.embedding = EmbeddingLayer(features + history_features + target_features)
+        self.interest_extractor_layers = nn.ModuleList([nn.GRU(fea.embed_dim, fea.embed_dim, batch_first=True) for fea in self.history_features])
+        self.interest_evolving_layers = nn.ModuleList([AUGRU(fea.embed_dim) for fea in self.history_features])
+
+        self.mlp = MLP(self.all_dims, activation="dice", **mlp_params)
+        self.history_labels = torch.Tensor(history_labels)
+        self.BCELoss = nn.BCELoss()
+
+
+# # 注意力计算中的线性层
+# self.attention_liner = nn.Linear(self.embed_dim, t)
+# # AFM公式中的h
+# self.h = init.xavier_uniform_(Parameter(torch.empty(t, 1)))
+# # AFM公式中的p
+# self.p = init.xavier_uniform_(Parameter(torch.empty(self.embed_dim, 1)))
+
+    def auxiliary(self, outs, history_features, history_labels):
+        '''
+        :param history_features: 历史序列物品的向量 [ batch_size, len_seqs, dim ]
+        :param outs: 兴趣抽取层GRU网络输出的outs [ batch_size, len_seqs, dim ]
+        :param history_labels: 历史序列物品标注 [ batch_size, len_seqs, 1 ]
+        :return: 辅助损失函数
+        '''
+        # [ batch_size * len_seqs, dim ]
+        history_features = history_features.reshape(-1, history_features.shape[2])
+        # [ batch_size * len_seqs, dim ]
+        outs = outs.reshape(-1, outs.shape[2])
+        # [ batch_size * len_seqs ]
+        out = torch.sum(outs * history_features, dim=1)
+        # [ batch_size * len_seqs, 1 ]
+        out = torch.unsqueeze(torch.sigmoid(out), 1)
+        # [ batch_size * len_seqs,1 ]
+        history_labels = history_labels.reshape(-1, 1).float()
+        return self.BCELoss(out, history_labels)
+
+    def forward(self, x):
+        # (batch_size, num_features, emb_dim)
+        embed_x_features = self.embedding(x, self.features)
+        # (batch_size, num_history_features, seq_length, emb_dim)
+        embed_x_history = self.embedding(x, self.history_features)
+        # (batch_size, num_target_features, emb_dim)
+        embed_x_target = self.embedding(x, self.target_features)
+
+        interest_extractor = []
+        auxi_loss = 0
+        for i in range(self.num_history_features):
+            outs, _ = self.interest_extractor_layers[i](embed_x_history[:, i, :, :])
+            # 利用GRU输出的outs得到辅助损失函数
+            auxi_loss += self.auxiliary(outs, embed_x_history[:, i, :, :], self.history_labels)
+            # (batch_size, 1, seq_length, emb_dim)
+            interest_extractor.append(outs.unsqueeze(1))
+        # (batch_size, num_history_features, seq_length, emb_dim)
+        interest_extractor = torch.cat(interest_extractor, dim=1)
+        interest_evolving = []
+        for i in range(self.num_history_features):
+            _, h = self.interest_evolving_layers[i](interest_extractor[:, i, :, :], embed_x_target[:, i, :])
+            interest_evolving.append(h.unsqueeze(1))  # (batch_size, 1, emb_dim)
+        # (batch_size, num_history_features, emb_dim)
+        interest_evolving = torch.cat(interest_evolving, dim=1)
+
+        mlp_in = torch.cat([interest_evolving.flatten(start_dim=1), embed_x_target.flatten(start_dim=1), embed_x_features.flatten(start_dim=1)], dim=1)  # (batch_size, N)
+        y = self.mlp(mlp_in)
+
+        return torch.sigmoid(y.squeeze(1)), self.alpha * auxi_loss

torch_rechub/models/ranking/din.py

@@ -1,7 +1,7 @@
 """
 Date: create on 23/04/2022, update on 30/04/2022
-References: 
-    paper: (KDD'2018) Deep Interest Network for Click-Through Rate Prediction 
+References:
+    paper: (KDD'2018) Deep Interest Network for Click-Through Rate Prediction
     url: https://arxiv.org/abs/1706.06978
     code: https://github.com/huawei-noah/benchmark/blob/main/FuxiCTR/fuxictr/pytorch/models/DIN.py
 Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
@@ -10,7 +10,7 @@ Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
 import torch
 import torch.nn as nn
 
-from ...basic.layers import EmbeddingLayer, MLP
+from ...basic.layers import MLP, EmbeddingLayer
 
 
 class DIN(nn.Module):
@@ -23,7 +23,7 @@ class DIN(nn.Module):
         attention_mlp_params (dict): the params of the ActivationUnit module, keys include:`{"dims":list, "activation":str, "dropout":float, "use_softmax":bool`}
     """
 
-    def __init__(self, features, history_features, target_features, mlp_params={"dims": [200, 80], "activation": "dice"}, attention_mlp_params={"dims": [36], "activation": "dice"}):
+    def __init__(self, features, history_features, target_features, mlp_params, attention_mlp_params):
         super().__init__()
         self.features = features
         self.history_features = history_features
@@ -33,19 +33,23 @@ class DIN(nn.Module):
 
         self.embedding = EmbeddingLayer(features + history_features + target_features)
         self.attention_layers = nn.ModuleList([ActivationUnit(fea.embed_dim, **attention_mlp_params) for fea in self.history_features])
-        self.mlp = MLP(self.all_dims, **mlp_params)
+        self.mlp = MLP(self.all_dims, activation="dice", **mlp_params)
 
     def forward(self, x):
-        embed_x_features = self.embedding(x, self.features)  #(batch_size, num_features, emb_dim)
-        embed_x_history = self.embedding(x, self.history_features)  #(batch_size, num_history_features, seq_length, emb_dim)
-        embed_x_target = self.embedding(x, self.target_features)  #(batch_size, num_target_features, emb_dim)
+        # (batch_size, num_features, emb_dim)
+        embed_x_features = self.embedding(x, self.features)
+        # (batch_size, num_history_features, seq_length, emb_dim)
+        embed_x_history = self.embedding(x, self.history_features)
+        # (batch_size, num_target_features, emb_dim)
+        embed_x_target = self.embedding(x, self.target_features)
         attention_pooling = []
         for i in range(self.num_history_features):
             attention_seq = self.attention_layers[i](embed_x_history[:, i, :, :], embed_x_target[:, i, :])
-            attention_pooling.append(attention_seq.unsqueeze(1))  #(batch_size, 1, emb_dim)
-        attention_pooling = torch.cat(attention_pooling, dim=1)  #(batch_size, num_history_features, emb_dim)
+            attention_pooling.append(attention_seq.unsqueeze(1))  # (batch_size, 1, emb_dim)
+        # (batch_size, num_history_features, emb_dim)
+        attention_pooling = torch.cat(attention_pooling, dim=1)
 
-        mlp_in = torch.cat([attention_pooling.flatten(start_dim=1), embed_x_target.flatten(start_dim=1), embed_x_features.flatten(start_dim=1)], dim=1)  #(batch_size, N)
+        mlp_in = torch.cat([attention_pooling.flatten(start_dim=1), embed_x_target.flatten(start_dim=1), embed_x_features.flatten(start_dim=1)], dim=1)  # (batch_size, N)
 
         y = self.mlp(mlp_in)
         return torch.sigmoid(y.squeeze(1))
@@ -62,20 +66,28 @@ class ActivationUnit(nn.Module):
         - Output: `(batch_size, emb_dim)`
     """
 
-    def __init__(self, emb_dim, dims=[36], activation="dice", use_softmax=False):
+    def __init__(self, emb_dim, dims=None, activation="dice", use_softmax=False):
         super(ActivationUnit, self).__init__()
+        if dims is None:
+            dims = [36]
         self.emb_dim = emb_dim
         self.use_softmax = use_softmax
-        self.attention = MLP(4 * self.emb_dim, dims, activation=activation)
+        self.attention = MLP(4 * self.emb_dim, dims=dims, activation=activation)
 
     def forward(self, history, target):
         seq_length = history.size(1)
-        target = target.unsqueeze(1).expand(-1, seq_length, -1)  #batch_size,seq_length,emb_dim
+        # batch_size,seq_length,emb_dim
+        target = target.unsqueeze(1).expand(-1, seq_length, -1)
         att_input = torch.cat([target, history, target - history, target * history], dim=-1)  # batch_size,seq_length,4*emb_dim
-        att_weight = self.attention(att_input.view(-1, 4 * self.emb_dim))  #  #(batch_size*seq_length,4*emb_dim)
-        att_weight = att_weight.view(-1, seq_length)  #(batch_size*seq_length, 1) -> (batch_size,seq_length)
+        # (batch_size*seq_length,4*emb_dim)
+        att_weight = self.attention(att_input.view(-1, 4 * self.emb_dim))
+        # (batch_size*seq_length, 1) -> (batch_size,seq_length)
+        att_weight = att_weight.view(-1, seq_length)
         if self.use_softmax:
             att_weight = att_weight.softmax(dim=-1)
-        # (batch_size, seq_length, 1) * (batch_size, seq_length, emb_dim)
-        output = (att_weight.unsqueeze(-1) * history).sum(dim=1)  #(batch_size,emb_dim)
-        return output
+
+
+# (batch_size, seq_length, 1) * (batch_size, seq_length, emb_dim)
+# (batch_size,emb_dim)
+        output = (att_weight.unsqueeze(-1) * history).sum(dim=1)
+        return output

torch_rechub/models/ranking/edcn.py

@@ -0,0 +1,101 @@
+"""
+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,42 @@
+"""
+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.features import SparseFeature
+from ...basic.layers import MLP, BiLinearInteractionLayer, EmbeddingLayer, SENETLayer
+
+
+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 is 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,7 +1,7 @@
 """
 Date: create on 22/04/2022
-References: 
-    paper: (DLRS'2016) Wide & Deep Learning for Recommender Systems 
+References:
+    paper: (DLRS'2016) Wide & Deep Learning for Recommender Systems
     url: https://arxiv.org/abs/1606.07792
 Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
 """
@@ -31,11 +31,11 @@ class WideDeep(torch.nn.Module):
         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]
+        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_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,4 @@
-from .trainer import CTRTrainer
-from .mtl_trainer import MTLTrainer
+from .ctr_trainer import CTRTrainer
+from .match_trainer import MatchTrainer
+from .mtl_trainer import MTLTrainer
+from .seq_trainer import SeqTrainer

torch_rechub/trainers/ctr_trainer.py

@@ -0,0 +1,191 @@
+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)