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

torch_rechub/models/multi_task/ple.py

@@ -1,6 +1,6 @@
 """
 Date: create on 05/05/2022
-References: 
+References:
     paper: (RecSys'2020) Progressive Layered Extraction (PLE): A Novel Multi-Task Learning (MTL) Model for Personalized Recommendations
     url: https://dl.acm.org/doi/abs/10.1145/3383313.3412236
 Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
@@ -21,11 +21,11 @@ class PLE(nn.Module):
         n_level (int): the  number of CGC layer.
         n_expert_specific (int): the number of task-specific expert net.
         n_expert_shared (int): the number of task-shared expert net.
-        expert_params (dict): the params of all the expert module, keys include:`{"dims":list, "activation":str, "dropout":float}, keep `{"output_layer":False}`.
+        expert_params (dict): the params of all the expert module, keys include:`{"dims":list, "activation":str, "dropout":float}.
         tower_params_list (list): the list of tower params dict, the keys same as expert_params.
     """
 
-    def __init__(self, features, task_types, n_level, n_expert_specific, n_expert_shared, expert_params={"dims": [32, 16], "output_layer": False}, tower_params_list=[{"dims": [32, 16]}, {"dims": [32, 16]}]):
+    def __init__(self, features, task_types, n_level, n_expert_specific, n_expert_shared, expert_params, tower_params_list):
         super().__init__()
         self.features = features
         self.n_task = len(task_types)
@@ -34,21 +34,25 @@ class PLE(nn.Module):
         self.input_dims = sum([fea.embed_dim for fea in features])
         self.embedding = EmbeddingLayer(features)
         self.cgc_layers = nn.ModuleList(CGC(i + 1, n_level, self.n_task, n_expert_specific, n_expert_shared, self.input_dims, expert_params) for i in range(n_level))
-        self.towers = nn.ModuleList(MLP(expert_params["dims"][-1], **tower_params_list[i]) for i in range(self.n_task))
+        self.towers = nn.ModuleList(MLP(expert_params["dims"][-1], output_layer=False, **tower_params_list[i]) for i in range(self.n_task))
         self.predict_layers = nn.ModuleList(PredictionLayer(task_type) for task_type in task_types)
 
     def forward(self, x):
-        embed_x = self.embedding(x, self.features, squeeze_dim=True)  #[batch_size, input_dims]
+        # [batch_size, input_dims]
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)
         ple_inputs = [embed_x] * (self.n_task + 1)
         ple_outs = []
         for i in range(self.n_level):
-            ple_outs = self.cgc_layers[i](ple_inputs)  #ple_outs[i]: [batch_size, expert_dims[-1]]
+            # ple_outs[i]: [batch_size, expert_dims[-1]]
+            ple_outs = self.cgc_layers[i](ple_inputs)
             ple_inputs = ple_outs
-        #predict
+
+
+# predict
         ys = []
         for ple_out, tower, predict_layer in zip(ple_outs, self.towers, self.predict_layers):
-            tower_out = tower(ple_out)  #[batch_size, 1]
-            y = predict_layer(tower_out)  #logit -> proba
+            tower_out = tower(ple_out)  # [batch_size, 1]
+            y = predict_layer(tower_out)  # logit -> proba
             ys.append(y)
         return torch.cat(ys, dim=1)
 
@@ -63,10 +67,10 @@ class CGC(nn.Module):
         n_expert_specific (int): the number of task-specific expert net.
         n_expert_shared (int): the number of task-shared expert net.
         input_dims (int): the input dims of the xpert module in current CGC layer.
-        expert_params (dict): the params of all the expert module, keys include:`{"dims":list, "activation":str, "dropout":float}, keep `{"output_layer":False}`.
+        expert_params (dict): the params of all the expert module, keys include:`{"dims":list, "activation":str, "dropout":float}.
     """
 
-    def __init__(self, cur_level, n_level, n_task, n_expert_specific, n_expert_shared, input_dims, expert_params={"dims": [32, 16], "output_layer": False}):
+    def __init__(self, cur_level, n_level, n_task, n_expert_specific, n_expert_shared, input_dims, expert_params):
         super().__init__()
         self.cur_level = cur_level  # the CGC level of PLE
         self.n_level = n_level
@@ -74,31 +78,39 @@ class CGC(nn.Module):
         self.n_expert_specific = n_expert_specific
         self.n_expert_shared = n_expert_shared
         self.n_expert_all = n_expert_specific * self.n_task + n_expert_shared
-        input_dims = input_dims if cur_level == 1 else expert_params["dims"][-1]  #the first layer expert dim is the input data dim other expert dim
-        self.experts_specific = nn.ModuleList(MLP(input_dims, **expert_params) for _ in range(self.n_task * self.n_expert_specific))
-        self.experts_shared = nn.ModuleList(MLP(input_dims, **expert_params) for _ in range(self.n_expert_shared))
-        self.gates_specific = nn.ModuleList(MLP(input_dims, **{"dims": [self.n_expert_specific + self.n_expert_shared], "activation": "softmax", "output_layer": False}) for _ in range(self.n_task))  #n_gate_specific = n_task
+        # the first layer expert dim is the input data dim other expert dim
+        input_dims = input_dims if cur_level == 1 else expert_params["dims"][-1]
+        self.experts_specific = nn.ModuleList(MLP(input_dims, output_layer=False, **expert_params) for _ in range(self.n_task * self.n_expert_specific))
+        self.experts_shared = nn.ModuleList(MLP(input_dims, output_layer=False, **expert_params) for _ in range(self.n_expert_shared))
+        self.gates_specific = nn.ModuleList(MLP(input_dims, **{"dims": [self.n_expert_specific + self.n_expert_shared], "activation": "softmax", "output_layer": False}) for _ in range(self.n_task))  # n_gate_specific = n_task
         if cur_level < n_level:
-            self.gate_shared = MLP(input_dims, **{"dims": [self.n_expert_all], "activation": "softmax", "output_layer": False})  #n_gate_specific = n_task
+            self.gate_shared = MLP(input_dims, **{"dims": [self.n_expert_all], "activation": "softmax", "output_layer": False})  # n_gate_specific = n_task
 
     def forward(self, x_list):
-        expert_specific_outs = []  #expert_out[i]: [batch_size, 1, expert_dims[-1]]
+        expert_specific_outs = []  # expert_out[i]: [batch_size, 1, expert_dims[-1]]
         for i in range(self.n_task):
             expert_specific_outs.extend([expert(x_list[i]).unsqueeze(1) for expert in self.experts_specific[i * self.n_expert_specific:(i + 1) * self.n_expert_specific]])
-        expert_shared_outs = [expert(x_list[-1]).unsqueeze(1) for expert in self.experts_shared]  #x_list[-1]: the input for shared experts
-        gate_specific_outs = [gate(x_list[i]).unsqueeze(-1) for i, gate in enumerate(self.gates_specific)]  #gate_out[i]: [batch_size, n_expert_specific+n_expert_shared, 1]
+        # x_list[-1]: the input for shared experts
+        expert_shared_outs = [expert(x_list[-1]).unsqueeze(1) for expert in self.experts_shared]
+        # gate_out[i]: [batch_size, n_expert_specific+n_expert_shared, 1]
+        gate_specific_outs = [gate(x_list[i]).unsqueeze(-1) for i, gate in enumerate(self.gates_specific)]
         cgc_outs = []
         for i, gate_out in enumerate(gate_specific_outs):
             cur_expert_list = expert_specific_outs[i * self.n_expert_specific:(i + 1) * self.n_expert_specific] + expert_shared_outs
-            expert_concat = torch.cat(cur_expert_list, dim=1)  #[batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
-            expert_weight = torch.mul(gate_out, expert_concat)  #[batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
-            expert_pooling = torch.sum(expert_weight, dim=1)  #[batch_size, expert_dims[-1]]
-            cgc_outs.append(expert_pooling)  #length: n_task
-        if self.cur_level < self.n_level:  #not the last layer
-            gate_shared_out = self.gate_shared(x_list[-1]).unsqueeze(-1)  #[batch_size, n_expert_all, 1]
-            expert_concat = torch.cat(expert_specific_outs + expert_shared_outs, dim=1)  #[batch_size, n_expert_all, expert_dims[-1]]
-            expert_weight = torch.mul(gate_shared_out, expert_concat)  #[batch_size, n_expert_all, expert_dims[-1]]
-            expert_pooling = torch.sum(expert_weight, dim=1)  #[batch_size, expert_dims[-1]]
-            cgc_outs.append(expert_pooling)  #length: n_task+1
+            # [batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
+            expert_concat = torch.cat(cur_expert_list, dim=1)
+            # [batch_size, n_expert_specific+n_expert_shared, expert_dims[-1]]
+            expert_weight = torch.mul(gate_out, expert_concat)
+            # [batch_size, expert_dims[-1]]
+            expert_pooling = torch.sum(expert_weight, dim=1)
+            cgc_outs.append(expert_pooling)  # length: n_task
+        if self.cur_level < self.n_level:  # not the last layer
+            gate_shared_out = self.gate_shared(x_list[-1]).unsqueeze(-1)  # [batch_size, n_expert_all, 1]
+            expert_concat = torch.cat(expert_specific_outs + expert_shared_outs, dim=1)  # [batch_size, n_expert_all, expert_dims[-1]]
+            # [batch_size, n_expert_all, expert_dims[-1]]
+            expert_weight = torch.mul(gate_shared_out, expert_concat)
+            # [batch_size, expert_dims[-1]]
+            expert_pooling = torch.sum(expert_weight, dim=1)
+            cgc_outs.append(expert_pooling)  # length: n_task+1
 
         return cgc_outs

torch_rechub/models/multi_task/shared_bottom.py

@@ -1,6 +1,6 @@
 """
 Date: create on 04/05/2022
-Reference: 
+Reference:
     paper: Caruana, R. (1997). Multitask learning. Machine learning, 28(1), 41-75.
 Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
 """
@@ -39,6 +39,7 @@ class SharedBottom(nn.Module):
         ys = []
         for tower, predict_layer in zip(self.towers, self.predict_layers):
             tower_out = tower(x)
-            y = predict_layer(tower_out)  #regression->keep, binary classification->sigmoid
+            # regression->keep, binary classification->sigmoid
+            y = predict_layer(tower_out)
             ys.append(y)
         return torch.cat(ys, dim=1)

torch_rechub/models/ranking/__init__.py

@@ -1,3 +1,14 @@
-from .widedeep import WideDeep
+__all__ = ['WideDeep', 'DeepFM', 'DCN', 'DCNv2', 'EDCN', 'AFM', 'FiBiNet', 'DeepFFM', 'BST', 'DIN', 'DIEN', 'FatDeepFFM', 'AutoInt']
+
+from .afm import AFM
+from .autoint import AutoInt
+from .bst import BST
+from .dcn import DCN
+from .dcn_v2 import DCNv2
+from .deepffm import DeepFFM, FatDeepFFM
 from .deepfm import DeepFM
-from .din import DIN
+from .dien import DIEN
+from .din import DIN
+from .edcn import EDCN
+from .fibinet import FiBiNet
+from .widedeep import WideDeep

torch_rechub/models/ranking/afm.py

@@ -0,0 +1,65 @@
+"""
+Date: create on 23/04/2024
+References:
+    paper: (IJCAI'2017) Attentional Factorization Machines：Learning the Weight of Feature Interactions via Attention Networks
+    url: https://arxiv.org/abs/1708.04617
+Authors: Tao Fan, thisisevy@foxmail.com
+"""
+
+import torch
+from torch import nn
+from torch.nn import Parameter, init
+
+from ...basic.layers import FM, LR, MLP, EmbeddingLayer
+
+
+class AFM(nn.Module):
+    """Attentional Factorization Machine Model
+
+    Args:
+        fm_features (list): the list of `Feature Class`, training by the fm part module.
+        embed_dim (int): the dimension of input embedding.
+        t (int): the size of the hidden layer in the attention network.
+    """
+
+    def __init__(self, fm_features, embed_dim, t=64):
+        super(AFM, self).__init__()
+        self.fm_features = fm_features
+        self.embed_dim = embed_dim
+        self.fm_dims = sum([fea.embed_dim for fea in fm_features])
+        self.linear = LR(self.fm_dims)  # 1-odrder interaction
+        self.fm = FM(reduce_sum=False)  # 2-odrder interaction
+        self.embedding = EmbeddingLayer(fm_features)
+
+        # 注意力计算中的线性层
+        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 attention(self, y_fm):
+        # embs: [ batch_size, k ]
+        # [ batch_size, t ]
+        y_fm = self.attention_liner(y_fm)
+        # [ batch_size, t ]
+        y_fm = torch.relu(y_fm)
+        # [ batch_size, 1 ]
+        y_fm = torch.matmul(y_fm, self.h)
+        # [ batch_size, 1 ]
+        atts = torch.softmax(y_fm, dim=1)
+        return atts
+
+    def forward(self, x):
+        # [batch_size, num_fields, embed_dim]
+        input_fm = self.embedding(x, self.fm_features, squeeze_dim=False)
+
+        y_linear = self.linear(input_fm.flatten(start_dim=1))
+        y_fm = self.fm(input_fm)
+        # 得到注意力
+        atts = self.attention(y_fm)
+        # [ batch_size, 1 ]
+        outs = torch.matmul(atts * y_fm, self.p)
+        # print(y_linear.size(), outs.size())
+        y = y_linear + outs
+        return torch.sigmoid(y.squeeze(1))

torch_rechub/models/ranking/autoint.py

@@ -0,0 +1,102 @@
+"""
+Date: create on 14/11/2025
+References:
+    paper: (CIKM'2019) AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks
+    url: https://arxiv.org/abs/1810.11921
+Authors: Yang Zhou, zyaztec@gmail.com
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import LR, MLP, EmbeddingLayer, InteractingLayer
+
+
+class AutoInt(torch.nn.Module):
+    """AutoInt Model
+
+    Args:
+        sparse_features (list): the list of `SparseFeature` Class
+        dense_features (list): the list of `DenseFeature` Class
+        num_layers (int): number of interacting layers
+        num_heads (int): number of attention heads
+        dropout (float): dropout rate for attention
+        mlp_params (dict): parameters for MLP, keys: {"dims":list, "activation":str,
+                                             "dropout":float, "output_layer":bool"}
+    """
+
+    def __init__(self, sparse_features, dense_features, num_layers=3, num_heads=2, dropout=0.0, mlp_params=None):
+        super(AutoInt, self).__init__()
+        self.sparse_features = sparse_features
+
+        self.dense_features = dense_features if dense_features is not None else []
+        embed_dims = [fea.embed_dim for fea in self.sparse_features]
+        self.embed_dim = embed_dims[0]
+        if len(self.sparse_features) == 0:
+            raise ValueError("AutoInt requires at least one sparse feature to determine embed_dim.")
+
+        # field nums = sparse + dense
+        self.num_sparse = len(self.sparse_features)
+        self.num_dense = len(self.dense_features)
+        self.num_fields = self.num_sparse + self.num_dense
+
+        # total dims = num_fields * embed_dim
+        self.dims = self.num_fields * self.embed_dim
+        self.num_layers = num_layers
+
+        self.sparse_embedding = EmbeddingLayer(self.sparse_features)
+
+        # dense feature embedding
+        self.dense_embeddings = nn.ModuleDict()
+        for fea in self.dense_features:
+            self.dense_embeddings[fea.name] = nn.Linear(1, self.embed_dim, bias=False)
+
+        self.interacting_layers = torch.nn.ModuleList([InteractingLayer(self.embed_dim, num_heads=num_heads, dropout=dropout, residual=True) for _ in range(num_layers)])
+
+        self.linear = LR(self.dims)
+
+        self.attn_linear = nn.Linear(self.dims, 1)
+
+        if mlp_params is not None:
+            self.use_mlp = True
+            self.mlp = MLP(self.dims, **mlp_params)
+        else:
+            self.use_mlp = False
+
+    def forward(self, x):
+        # sparse feature embedding: [B, num_sparse, embed_dim]
+        sparse_emb = self.sparse_embedding(x, self.sparse_features, squeeze_dim=False)
+
+        dense_emb_list = []
+        for fea in self.dense_features:
+            v = x[fea.name].float().view(-1, 1, 1)
+            dense_emb = self.dense_embeddings[fea.name](v)  # [B, 1, embed_dim]
+            dense_emb_list.append(dense_emb)
+
+        if len(dense_emb_list) > 0:
+            dense_emb = torch.cat(dense_emb_list, dim=1)  # [B, num_dense, d]
+            embed_x = torch.cat([sparse_emb, dense_emb], dim=1)  # [B, num_fields, d]
+        else:
+            embed_x = sparse_emb  # [B, num_sparse, d]
+
+        embed_x_flatten = embed_x.flatten(start_dim=1)  # [B, num_fields * embed_dim]
+
+        # Multi-head self-attention layers
+        attn_out = embed_x
+        for layer in self.interacting_layers:
+            attn_out = layer(attn_out)  # [B, num_fields, embed_dim]
+
+        # Attention linear
+        attn_out_flatten = attn_out.flatten(start_dim=1)  # [B, num_fields * embed_dim]
+        y_attn = self.attn_linear(attn_out_flatten)  # [B, 1]
+
+        # Linear part
+        y_linear = self.linear(embed_x_flatten)  # [B, 1]
+
+        # Deep MLP
+        y = y_attn + y_linear
+        if self.use_mlp:
+            y_deep = self.mlp(embed_x_flatten)  # [B, 1]
+            y = y + y_deep
+
+        return torch.sigmoid(y.squeeze(1))

torch_rechub/models/ranking/bst.py

@@ -0,0 +1,61 @@
+"""
+Date: create on 26/02/2024, update on 30/04/2022
+References:
+    paper: Behavior Sequence Transformer for E-commerce Recommendation in Alibaba
+    url: https://arxiv.org/pdf/1905.06874
+    code: https://github.com/jiwidi/Behavior-Sequence-Transformer-Pytorch/blob/master/pytorch_bst.ipynb
+Authors: Tao Fan, thisisevy@foxmail.com
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import MLP, EmbeddingLayer
+
+
+class BST(nn.Module):
+    """Behavior Sequence Transformer
+    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`}.
+        nhead (int): the number of heads in the multi-head-attention models.
+        dropout (float): the dropout value in the multi-head-attention models.
+        num_layers (Any): the number of sub-encoder-layers in the encoder.
+    """
+
+    def __init__(self, features, history_features, target_features, mlp_params, nhead=8, dropout=0.2, num_layers=1):
+        super().__init__()
+        self.features = features
+        self.history_features = history_features
+        self.target_features = target_features
+        self.num_history_features = len(history_features)
+        self.embed_dim = target_features[0].embed_dim
+        self.seq_len = 50
+        # TODO 在 'torch_rechub.basic.features.SequenceFeature' 中加入seq_len属性
+        self.all_dims = (len(features) + len(history_features) * (self.seq_len + len(target_features))) * self.embed_dim
+        self.embedding = EmbeddingLayer(features + history_features + target_features)
+        self.encoder_layer = nn.TransformerEncoderLayer(d_model=self.embed_dim, nhead=nhead, dropout=dropout)
+        self.transformer_layers = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
+        # # 定义模型，模型的参数需要我们之前的feature类，用于构建模型的输入层，mlp指定模型后续DNN的结构
+        self.mlp = MLP(self.all_dims, activation="leakyrelu", **mlp_params)
+
+    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)
+        attention_pooling = []
+        for i in range(self.num_history_features):
+            attention_seq = self.transformer_layers(torch.cat([torch.squeeze(embed_x_history[:, i, :, :], 1), embed_x_target], dim=1))
+            # (batch_size, seq_length + num_target_features, emb_dim)
+            attention_pooling.append(attention_seq)
+        # (batch_size, num_history_features * (seq_length + num_target_features), emb_dim)
+        attention_pooling = torch.cat(attention_pooling, dim=1)
+
+        mlp_in = torch.cat([attention_pooling.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))

torch_rechub/models/ranking/dcn.py

@@ -0,0 +1,38 @@
+"""
+Date: create on 12/05/2022
+References:
+    paper: (AKDD'2017) Deep & Cross Network for Ad Click Predictions
+    url: https://arxiv.org/abs/1708.05123
+Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
+"""
+
+import torch
+
+from ...basic.layers import LR, MLP, CrossNetwork, EmbeddingLayer
+
+
+class DCN(torch.nn.Module):
+    """Deep & Cross Network
+
+    Args:
+        features (list[Feature Class]): training by the whole 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, features, n_cross_layers, mlp_params):
+        super().__init__()
+        self.features = features
+        self.dims = sum([fea.embed_dim for fea in features])
+
+        self.embedding = EmbeddingLayer(features)
+        self.cn = CrossNetwork(self.dims, n_cross_layers)
+        self.mlp = MLP(self.dims, output_layer=False, **mlp_params)
+        self.linear = LR(self.dims + mlp_params["dims"][-1])
+
+    def forward(self, x):
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)
+        cn_out = self.cn(embed_x)
+        mlp_out = self.mlp(embed_x)
+        x_stack = torch.cat([cn_out, mlp_out], dim=1)
+        y = self.linear(x_stack)
+        return torch.sigmoid(y.squeeze(1))

torch_rechub/models/ranking/dcn_v2.py

@@ -0,0 +1,59 @@
+"""
+Date: create on 09/01/2022
+References:
+    paper: (WWW'21) Dcn v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems
+    url: https://arxiv.org/abs/2008.13535
+Authors: lailai, lailai_zxy@tju.edu.cn
+"""
+import torch
+
+from ...basic.layers import LR, MLP, CrossNetMix, CrossNetV2, EmbeddingLayer
+
+
+class DCNv2(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`}
+        use_low_rank_mixture (bool): True, whether to use a mixture of low-rank architecture
+        low_rank (int): the rank size of low-rank matrices
+        num_experts (int): the number of expert networks
+    """
+
+    def __init__(self, features, n_cross_layers, mlp_params, model_structure="parallel", use_low_rank_mixture=True, low_rank=32, num_experts=4, **kwargs):
+        super(DCNv2, self).__init__()
+        self.features = features
+        self.dims = sum([fea.embed_dim for fea in features])
+        self.embedding = EmbeddingLayer(features)
+        if use_low_rank_mixture:
+            self.crossnet = CrossNetMix(self.dims, n_cross_layers, low_rank=low_rank, num_experts=num_experts)
+        else:
+            self.crossnet = CrossNetV2(self.dims, n_cross_layers)
+        self.model_structure = model_structure
+        assert self.model_structure in ["crossnet_only", "stacked", "parallel"], \
+            "model_structure={} not supported!".format(self.model_structure)
+        if self.model_structure == "stacked":
+            self.stacked_dnn = MLP(self.dims, output_layer=False, **mlp_params)
+            final_dim = mlp_params["dims"][-1]
+        if self.model_structure == "parallel":
+            self.parallel_dnn = MLP(self.dims, output_layer=False, **mlp_params)
+            final_dim = mlp_params["dims"][-1] + self.dims
+        if self.model_structure == "crossnet_only":  # only CrossNet
+            final_dim = self.dims
+        self.linear = LR(final_dim)
+
+    def forward(self, x):
+        embed_x = self.embedding(x, self.features, squeeze_dim=True)
+        cross_out = self.crossnet(embed_x)
+        if self.model_structure == "crossnet_only":
+            final_out = cross_out
+        elif self.model_structure == "stacked":
+            final_out = self.stacked_dnn(cross_out)
+        elif self.model_structure == "parallel":
+            dnn_out = self.parallel_dnn(embed_x)
+            final_out = torch.cat([cross_out, dnn_out], dim=1)
+        y_pred = self.linear(final_out)
+        y_pred = torch.sigmoid(y_pred.squeeze(1))
+        return y_pred

torch_rechub/models/ranking/deepffm.py

@@ -0,0 +1,131 @@
+"""
+Date: created on 31/07/2022
+References:
+    paper: FAT-DeepFFM: Field Attentive Deep Field-aware Factorization Machine
+    url: https://arxiv.org/abs/1905.06336
+Authors: Bo Kang, klinux@live.com
+"""
+
+import torch
+import torch.nn as nn
+
+from ...basic.layers import CEN, FFM, MLP, EmbeddingLayer
+
+
+class DeepFFM(nn.Module):
+    """The DeepFFM model, mentioned on the `webpage
+    <https://cs.nju.edu.cn/31/60/c1654a209248/page.htm>` which is the first
+    work that introduces FFM model into neural CTR system. It is also described
+    in the `FAT-DeepFFM paper <https://arxiv.org/abs/1905.06336>`.
+
+    Args:
+        linear_features (list): the list of `Feature Class`, fed to the linear module.
+        cross_features (list): the list of `Feature Class`, fed to the ffm module.
+        embed_dim (int): the dimensionality of categorical value embedding.
+        mlp_params (dict): the params of the last MLP module, keys include:`{"dims":list, "activation":str, "dropout":float, "output_layer":bool`}
+    """
+
+    def __init__(self, linear_features, cross_features, embed_dim, mlp_params):
+        super().__init__()
+        self.linear_features = linear_features
+        self.cross_features = cross_features
+
+        self.num_fields = len(cross_features)
+        self.num_field_cross = self.num_fields * (self.num_fields - 1) // 2
+
+        self.ffm = FFM(num_fields=self.num_fields, reduce_sum=False)
+        self.mlp_out = MLP(self.num_field_cross * embed_dim, **mlp_params)
+
+        self.linear_embedding = EmbeddingLayer(linear_features)
+        self.ffm_embedding = EmbeddingLayer(cross_features)
+
+        self.b = torch.nn.Parameter(torch.zeros(1))
+
+        # This keeping constant value in module on correct device
+        # url:
+        # https://discuss.pytorch.org/t/keeping-constant-value-in-module-on-correct-device/10129
+        fields_offset = torch.arange(0, self.num_fields, dtype=torch.long)
+        self.register_buffer('fields_offset', fields_offset)
+
+    def forward(self, x):
+        # compute scores from the linear part of the model, where input is the
+        # raw features (Eq. 5, FAT-DeepFFM)
+        y_linear = self.linear_embedding(x, self.linear_features, squeeze_dim=True).sum(1, keepdim=True)  # [batch_size, 1]
+
+        # gather the embeddings. Each feature value corresponds to multiple embeddings, with multiplicity equal to number of features/fields.
+        # output shape [batch_size, num_field, num_field, emb_dim]
+        x_ffm = {fea.name: x[fea.name].unsqueeze(1) * self.num_fields + self.fields_offset for fea in self.cross_features}
+        input_ffm = self.ffm_embedding(x_ffm, self.cross_features, squeeze_dim=False)
+
+        # compute second order field-aware feature crossings, output shape
+        # [batch_size, num_field_cross, emb_dim]
+        em = self.ffm(input_ffm)
+
+        # compute scores from the ffm part of the model, output shape
+        # [batch_size, 1]
+        y_ffm = self.mlp_out(em.flatten(start_dim=1))
+
+        # compute final prediction
+        y = y_linear + y_ffm
+        return torch.sigmoid(y.squeeze(1) + self.b)
+
+
+class FatDeepFFM(nn.Module):
+    """The FAT-DeepFFM model, mentioned in the `FAT-DeepFFM paper
+    <https://arxiv.org/abs/1905.06336>`. It combines DeepFFM with
+    Compose-Excitation Network (CENet) field attention mechanism
+    to highlight the importance of second-order feature crosses.
+
+    Args:
+        linear_features (list): the list of `Feature Class`, fed to the linear module.
+        cross_features (list): the list of `Feature Class`, fed to the ffm module.
+        embed_dim (int): the dimensionality of categorical value embedding.
+        reduction_ratio (int): the between the dimensions of input layer and hidden layer of the CEN MLP 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, linear_features, cross_features, embed_dim, reduction_ratio, mlp_params):
+        super().__init__()
+        self.linear_features = linear_features
+        self.cross_features = cross_features
+
+        self.num_fields = len(cross_features)
+        self.num_field_cross = self.num_fields * (self.num_fields - 1) // 2
+
+        self.ffm = FFM(num_fields=self.num_fields, reduce_sum=False)
+        self.cen = CEN(embed_dim, self.num_field_cross, reduction_ratio)
+        self.mlp_out = MLP(self.num_field_cross * embed_dim, **mlp_params)
+
+        self.linear_embedding = EmbeddingLayer(linear_features)
+        self.ffm_embedding = EmbeddingLayer(cross_features)
+
+        self.b = torch.nn.Parameter(torch.zeros(1))
+
+        fields_offset = torch.arange(0, self.num_fields, dtype=torch.long)
+        self.register_buffer('fields_offset', fields_offset)
+
+    def forward(self, x):
+        # compute scores from the linear part of the model, where input is the
+        # raw features (Eq. 5, FAT-DeepFFM)
+        y_linear = self.linear_embedding(x, self.linear_features, squeeze_dim=True).sum(1, keepdim=True)  # [batch_size, 1]
+
+        # gather the embeddings. Each feature value corresponds to multiple embeddings, with multiplicity is equal to the number of features/fields.
+        # output shape [batch_size, num_field, num_field, emb_dim]
+        x_ffm = {fea.name: x[fea.name].unsqueeze(1) * self.num_fields + self.fields_offset for fea in self.cross_features}
+        input_ffm = self.ffm_embedding(x_ffm, self.cross_features, squeeze_dim=False)
+
+        # compute second order field-aware feature crossings, output shape
+        # [batch_size, num_field_cross, emb_dim]
+        em = self.ffm(input_ffm)
+
+        # rescale FFM embeddings with field attention (Eq.10), output shape
+        # [batch_size, num_field_cross * emb_dim]
+        aem = self.cen(em)
+
+        # compute scores from the ffm part of the model, output shape
+        # [batch_size, 1]
+        y_ffm = self.mlp_out(aem)
+
+        # compute final prediction
+        y = y_linear + y_ffm
+        return torch.sigmoid(y.squeeze(1) + self.b)

torch_rechub/models/ranking/deepfm.py

@@ -1,14 +1,14 @@
 """
 Date: create on 22/04/2022
-References: 
-    paper: (IJCAI'2017) DeepFM: A Factorization-Machine based Neural Network for CTR Prediction 
+References:
+    paper: (IJCAI'2017) DeepFM: A Factorization-Machine based Neural Network for CTR Prediction
     url: https://arxiv.org/abs/1703.04247
 Authors: Mincai Lai, laimincai@shanghaitech.edu.cn
 """
 
 import torch
 
-from ...basic.layers import FM, MLP, LR, EmbeddingLayer
+from ...basic.layers import FM, LR, MLP, EmbeddingLayer
 
 
 class DeepFM(torch.nn.Module):
@@ -32,11 +32,12 @@ class DeepFM(torch.nn.Module):
         self.mlp = MLP(self.deep_dims, **mlp_params)
 
     def forward(self, x):
-        input_deep = self.embedding(x, self.deep_features, squeeze_dim=True)  #[batch_size, deep_dims]
-        input_fm = self.embedding(x, self.fm_features, squeeze_dim=False)  #[batch_size, num_fields, embed_dim]
+        input_deep = self.embedding(x, self.deep_features, squeeze_dim=True)  # [batch_size, deep_dims]
+        # [batch_size, num_fields, embed_dim]
+        input_fm = self.embedding(x, self.fm_features, squeeze_dim=False)
 
         y_linear = self.linear(input_fm.flatten(start_dim=1))
         y_fm = self.fm(input_fm)
-        y_deep = self.mlp(input_deep)  #[batch_size, 1]
+        y_deep = self.mlp(input_deep)  # [batch_size, 1]
         y = y_linear + y_fm + y_deep
-        return torch.sigmoid(y.squeeze(1))
+        return torch.sigmoid(y.squeeze(1))