diff --git a/docs/images/models/pepnet.jpg b/docs/images/models/pepnet.jpg new file mode 100644 index 00000000..54580229 Binary files /dev/null and b/docs/images/models/pepnet.jpg differ diff --git a/docs/source/models/multi_target.rst b/docs/source/models/multi_target.rst index ea55acc3..59bd7958 100644 --- a/docs/source/models/multi_target.rst +++ b/docs/source/models/multi_target.rst @@ -7,3 +7,4 @@ mmoe dbmtl ple + pepnet diff --git a/docs/source/models/pepnet.md b/docs/source/models/pepnet.md new file mode 100644 index 00000000..617422c8 --- /dev/null +++ b/docs/source/models/pepnet.md @@ -0,0 +1,121 @@ +# PEPNet + +## 简介 + +PEPNet将带有个性化先验信息的特征作为输入，通过门控机制，来动态地缩放底层网络-embedding layer和顶层网络-DNN隐藏层单元，分别称之为场景特定的EPNet和任务特定的PPNet： + +- Embedding Personalized Network (EPNet) 为底层网络添加场景特定的个性化信息来生成个性化embedding门控，用于执行来自多个场景的原始embedding选择，以生成个性化的embedding +- Parameter Personalized Network (PPNet) 将用户和items的个性化信息与每一个task tower的DNN的输入进行拼接来获得个性化的门控分数，然后采用element-wise product应用到DNN的隐藏层单元上，来个性化优化DNN的参数。 + +PEPNet的整体结构如下图所示，可以看到，核心的组件便是这三个：Gate NU（门控网络单元）、Embedding Personalized Network (EPNet)、Parameter Personalized Network (PPNet) +![pepnet.jpg](../../images/models/pepnet.jpg) + +## 模型配置 + +```protobuf +model_config: { + feature_groups: { + group_name: 'all' + feature_names: 'user_id' + feature_names: 'cms_segid' + feature_names: 'cms_group_id' + feature_names: 'age_level' + feature_names: 'pvalue_level' + feature_names: 'shopping_level' + feature_names: 'occupation' + feature_names: 'new_user_class_level' + feature_names: 'adgroup_id' + feature_names: 'cate_id' + feature_names: 'campaign_id' + feature_names: 'customer' + feature_names: 'brand' + feature_names: 'pid' + feature_names: 'tag_category_list' + feature_names: 'tag_brand_list' + feature_names: 'price' + wide_deep: DEEP + } + feature_groups: { + group_name: 'domain' + feature_names: 'occupation' + wide_deep: DEEP + } + feature_groups: { + group_name: 'uia' + feature_names: 'user_id' + feature_names: 'adgroup_id' + wide_deep: DEEP + } + pepnet { + main_group_name: "all" + domain_group_name: "domain" + epnet_hidden_unit: 128, + uia_group_name: "uia" + ppnet_hidden_units: [512, 256] + ppnet_dropout_ratio: [0.1, 0.1] + task_towers { + tower_name: "ctr" + label_name: "clk" + mlp { + hidden_units: [256, 128, 64] + } + metrics { + auc {} + } + losses { + binary_cross_entropy {} + } + } + task_towers { + tower_name: "cvr" + label_name: "buy" + mlp { + hidden_units: [256, 128, 64] + } + metrics { + auc { + thresholds: 1000 + } + } + losses { + binary_cross_entropy {} + } + } + } +} +``` + +- feature_groups: 特征组 + + - 通常情况下有3个feature_group: 名称自定义，根据pepnet的配置，分为3类：all, domain, uia，其中domain和uia是可选配置，根据需求进行配置。 + - wide_deep: pepnet模型使用的都是Deep features, 所以都设置成DEEP + +- pepnet: pepnet模型相关的参数 + + - main_group_name: 主特征组名称,和feature_groups中的group_name对应 + - domain_group_name: domain特征组名称,和feature_groups中的group_name对应，是epnet场景个性化部分的输入 + - epnet_hidden_unit: epnet的gateGu的隐层设置，一般介于domain的dim和主特征组dim之间 + - uia_group_name: 用户和item的个性化信息，和feature_groups中的group_name对应，是ppnet用户-商品个性化部分的输入 + - ppnet_hidden_units: 个性化tower的隐藏层设置 + - task_towers: 根据任务数配置task_towers + - tower_name：TaskTower名 + - label_name: tower对应的label名 + - mlp: TaskTower的MLP参数配置 + - weight: 任务权重名 + - sample_weight_name: 样本权重列名 + - losses: 任务损失函数配置 + - metrics: 任务评估指标配置 + - task_space_indicator_label: 标识当前任务空间的目标名称，配合in_task_space_weight、out_task_space_weight使用。例如，对于cvr任务，可以设置task_space_indicator_label=clk，in_task_space_weight=1，out_task_space_weight=0，来使得cvr任务塔只在点击空间计算loss。 + - 注: in_task_space_weight和out_task_space_weight不影响loss权重的绝对值，权重会在batch维度被归一化。例如：in_task_space_weight=10,out_task_space_weight=1跟in_task_space_weight=1,out_task_space_weight=0.1是等价的。如需要提升这个task的loss权重的绝对值，需设置weight参数 + - in_task_space_weight: 对于task_space_indicator_label>0的样本会乘以该权重 + - out_task_space_weight: 对于task_space_indicator_label\<=0的样本会乘以该权重 + +## 模型输出 + +和其余多任务模型一样，每个塔的输出名为："logits\_" / "probs\_" / "y\_" + tower_name +其中，logits/probs/y对应: sigmoid之前的值/概率/回归模型的预测值 +MMoE模型每个塔的指标为：指标名+ "\_" + tower_name + +## 参考论文 + +[PEPNet.pdf](https://arxiv.org/pdf/2302.01115) diff --git a/tzrec/models/pepnet.py b/tzrec/models/pepnet.py new file mode 100644 index 00000000..1c1bd938 --- /dev/null +++ b/tzrec/models/pepnet.py @@ -0,0 +1,286 @@ +# Copyright (c) 2024, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict, List, Optional, Union + +import torch +from torch import nn + +from tzrec.datasets.utils import Batch +from tzrec.features.feature import BaseFeature +from tzrec.models.multi_task_rank import MultiTaskRank +from tzrec.modules.activation import create_activation +from tzrec.modules.task_tower import TaskTower +from tzrec.protos.model_pb2 import ModelConfig +from tzrec.utils.config_util import config_to_kwargs + + +class GateNU(nn.Module): + """Gate Neural Unit for PEPNet. + + Implements the Gate Neural Unit from the PEPNet paper with ReLU + -> Sigmoid activation and optional gamma scaling factor. + """ + + def __init__( + self, input_dim: int, hidden_dim: int, output_dim: int, gamma: float = 2.0 + ) -> None: + super().__init__() + self._gamma = gamma + self._output_dim = output_dim + self.dense_layers = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, output_dim), + nn.Sigmoid(), + ) + + def output_dim(self) -> int: + """Get output dimension of the GateNU.""" + return self._output_dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass for GateNU. + + Args: + x: Input tensor [B, input_dim] + + Returns: + Output tensor [B, output_dim] with values scaled by gamma + """ + return self._gamma * self.dense_layers(x) + + +class EPNet(nn.Module): + """Embedding Personalization Network for PEPNet. + + Generates personalized weights to scale the original embeddings + based on domain and context information, following the original PEPNet paper. + """ + + def __init__( + self, main_dim: int, domain_dim: int, hidden_dim: int, gamma: float = 2.0 + ) -> None: + super().__init__() + self._domain_dim = domain_dim + self._main_dim = main_dim + input_dim = domain_dim + main_dim + self.gate_nu = GateNU( + input_dim=input_dim, hidden_dim=hidden_dim, output_dim=main_dim, gamma=gamma + ) + + def output_dim(self) -> int: + """Get output dimension of the EPNet.""" + return self.gate_nu.output_dim() + + def forward( + self, + main_emb: torch.Tensor, + domain_emb: torch.Tensor, + ) -> torch.Tensor: + """Forward pass for embedding personalization. + + Args: + main_emb: Main feature embedding tensor [B, embedding_dim] + domain_emb: Domain embedding tensor [B, domain_dim] + + Returns: + Personalized embedding tensor [B, embedding_dim] + """ + gate_input = torch.cat([domain_emb, main_emb.detach()], dim=-1) + scaling_factors = self.gate_nu(gate_input) + personalized_emb = scaling_factors * main_emb + return personalized_emb + + +class PPNet(nn.Module): + """Parameter Personalization Network for PEPNet.""" + + def __init__( + self, + main_feature: int, + uia_feature: int, + num_task: int, + hidden_units: List[int], + activation: Optional[str] = "nn.ReLU", + dropout_ratio: Optional[Union[List[float], float]] = None, + gamma: float = 2.0, + ) -> None: + super().__init__() + self.main_feature = main_feature + self.uia_feature = uia_feature + self.num_task = num_task + self.hidden_units = hidden_units + self.len_hidden = len(hidden_units) + self.linears = nn.ModuleList() + self.activations = nn.ModuleList() + self.dropout_ratios = nn.ModuleList() + self.gate_nus = nn.ModuleList() + + if dropout_ratio is None: + dropout_ratio = [0.0] * len(hidden_units) + elif isinstance(dropout_ratio, list): + if len(dropout_ratio) == 0: + dropout_ratio = [0.0] * len(hidden_units) + elif len(dropout_ratio) == 1: + dropout_ratio = dropout_ratio * len(hidden_units) + else: + assert len(dropout_ratio) == len(hidden_units), ( + "length of dropout_ratio and hidden_units must be same, " + f"but got {len(dropout_ratio)} vs {len(hidden_units)}" + ) + else: + dropout_ratio = [dropout_ratio] * len(hidden_units) + + for _ in range(self.num_task): + output = main_feature + for i, hidden_unit in enumerate(hidden_units): + self.linears.append(nn.Linear(output, hidden_unit)) + active = create_activation(activation, hidden_size=hidden_unit, dim=2) + self.activations.append(active) + self.dropout_ratios.append(nn.Dropout(dropout_ratio[i])) + self.gate_nus.append( + GateNU( + input_dim=self.main_feature + self.uia_feature, + hidden_dim=hidden_unit, + output_dim=hidden_unit, + gamma=gamma, + ) + ) + output = hidden_unit + + def output_dim(self) -> List[int]: + """Get output dimension of the PPNet.""" + return [self.hidden_units[-1]] * self.num_task + + def task_output_dim(self) -> int: + """Get output dimension of the PPNet.""" + return self.hidden_units[-1] + + def forward( + self, main_emb: torch.Tensor, uia_emb: torch.Tensor + ) -> List[torch.Tensor]: + """Forward pass for parameter personalization. + + Args: + main_emb: Input tensor [B, embedding_dim] + uia_emb: UI personalization embedding tensor [B, uia_dim] + + Returns: + List of task-specific outputs [[B, output_dim]] * num_task + """ + task_outputs = [] + for i in range(self.num_task): + gate_input = torch.cat([uia_emb, main_emb.detach()], dim=-1) + input = main_emb + for j in range(self.len_hidden): + input = self.linears[i * self.len_hidden + j](input) + input = self.activations[i * self.len_hidden + j](input) + input = input * self.gate_nus[i * self.len_hidden + j](gate_input) + input = self.dropout_ratios[i * self.len_hidden + j](input) + task_outputs.append(input) + return task_outputs + + +class PEPNet(MultiTaskRank): + """Parameter and Embedding Personalized Network.""" + + def __init__( + self, + model_config: ModelConfig, + features: List[BaseFeature], + labels: List[str], + sample_weights: Optional[List[str]] = None, + **kwargs: Any, + ) -> None: + super().__init__(model_config, features, labels, sample_weights, **kwargs) + self.init_input() + + self.main_group_name = self._model_config.main_group_name + self.main_group_dim = self.embedding_group.group_total_dim(self.main_group_name) + self.task_input_dim = self.main_group_dim + self.domain_group_name = None + self.epnet = None + if self._model_config.HasField("domain_group_name"): + self.domain_group_name = self._model_config.domain_group_name + domain_group_dim = self.embedding_group.group_total_dim( + self.domain_group_name + ) + self.epnet = EPNet( + self.main_group_dim, + domain_group_dim, + hidden_dim=self._model_config.epnet_hidden_unit, + gamma=self._model_config.epnet_gamma, + ) + self.task_input_dim = self.epnet.output_dim() + + self.uia_group_name = None + self.ppnet = None + if self._model_config.HasField("uia_group_name"): + self.uia_group_name = self._model_config.uia_group_name + uia_group_dim = self.embedding_group.group_total_dim(self.uia_group_name) + self.ppnet = PPNet( + self.main_group_dim, + uia_group_dim, + num_task=len(self._task_tower_cfgs), + hidden_units=list(self._model_config.ppnet_hidden_units), + activation=self._model_config.ppnet_activation, + dropout_ratio=list(self._model_config.ppnet_dropout_ratio), + gamma=self._model_config.ppnet_gamma, + ) + self.task_input_dim = self.ppnet.task_output_dim() + + self._task_tower = nn.ModuleList() + for tower_cfg in self._task_tower_cfgs: + tower_cfg = config_to_kwargs(tower_cfg) + mlp = tower_cfg["mlp"] if "mlp" in tower_cfg else None + self._task_tower.append( + TaskTower(self.task_input_dim, tower_cfg["num_class"], mlp=mlp) + ) + + def predict(self, batch: Batch) -> Dict[str, torch.Tensor]: + """Forward the model. + + Args: + batch (Batch): input batch data. + + Return: + predictions (dict): a dict of predicted result. + """ + grouped_features = self.build_input(batch) + + # Get main features + main_features = grouped_features[self.main_group_name] + # Apply EPNet if available for embedding personalization + if self.domain_group_name: + domain_features = grouped_features[self.domain_group_name] + final_features = self.epnet(main_features, domain_features) + else: + final_features = main_features + + if self.uia_group_name: + uia_features = grouped_features[self.uia_group_name] + task_input_list = self.ppnet(final_features, uia_features) + else: + task_input_list = [final_features] + + # Apply task towers + tower_outputs = {} + for i, task_tower_cfg in enumerate(self._task_tower_cfgs): + tower_name = task_tower_cfg.tower_name + if self.uia_group_name: + task_input = task_input_list[i] + else: + task_input = task_input_list[0] + tower_output = self._task_tower[i](task_input) + tower_outputs[tower_name] = tower_output + + return self._multi_task_output_to_prediction(tower_outputs) diff --git a/tzrec/models/pepnet_test.py b/tzrec/models/pepnet_test.py new file mode 100644 index 00000000..66450701 --- /dev/null +++ b/tzrec/models/pepnet_test.py @@ -0,0 +1,455 @@ +# Copyright (c) 2024, Alibaba Group; +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch +from parameterized import parameterized +from torchrec import KeyedJaggedTensor, KeyedTensor + +from tzrec.datasets.utils import BASE_DATA_GROUP, Batch +from tzrec.features.feature import create_features +from tzrec.models.pepnet import EPNet, GateNU, PEPNet, PPNet +from tzrec.protos import ( + feature_pb2, + loss_pb2, + model_pb2, + module_pb2, + seq_encoder_pb2, + tower_pb2, +) +from tzrec.protos.models import multi_task_rank_pb2 +from tzrec.utils.state_dict_util import init_parameters +from tzrec.utils.test_util import TestGraphType, create_test_model, create_test_module + + +class GateNUTest(unittest.TestCase): + """Test GateNU module.""" + + @parameterized.expand( + [[TestGraphType.NORMAL], [TestGraphType.FX_TRACE], [TestGraphType.JIT_SCRIPT]] + ) + def test_gatenu(self, graph_type) -> None: + """Test GateNU forward pass.""" + input_dim = 32 + hidden_dim = 16 + output_dim = 8 + + gatenu = GateNU( + input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, gamma=2.0 + ) + gatenu = create_test_module(gatenu, graph_type) + + batch_size = 4 + x = torch.randn(batch_size, input_dim) + + output = gatenu(x) + + self.assertEqual(output.shape, (batch_size, output_dim)) + # Check that output is in [0, gamma] range due to sigmoid activation + self.assertTrue(torch.all(output >= 0)) + self.assertTrue(torch.all(output <= 2.0)) + + +class EPNetTest(unittest.TestCase): + """Test EPNet module.""" + + @parameterized.expand( + [[TestGraphType.NORMAL], [TestGraphType.FX_TRACE], [TestGraphType.JIT_SCRIPT]] + ) + def test_epnet(self, graph_type) -> None: + """Test EPNet forward pass.""" + domain_dim = 16 + embedding_dim = 32 + epnet = EPNet( + main_dim=embedding_dim, + domain_dim=domain_dim, + hidden_dim=8, + ) + epnet = create_test_module(epnet, graph_type) + batch_size = 4 + domain_emb = torch.randn(batch_size, domain_dim) + main_emb = torch.randn(batch_size, embedding_dim) + personalized_emb = epnet(main_emb, domain_emb) + self.assertEqual(personalized_emb.shape, (batch_size, embedding_dim)) + + +class PPNetTest(unittest.TestCase): + """Test PPNet module.""" + + @parameterized.expand( + [[TestGraphType.NORMAL], [TestGraphType.FX_TRACE], [TestGraphType.JIT_SCRIPT]] + ) + def test_ppnet_forward(self, graph_type) -> None: + """Test PPNet forward pass.""" + main_feature = 32 + uia_feature = 16 + num_task = 2 + + ppnet = PPNet( + main_feature=main_feature, + uia_feature=uia_feature, + num_task=num_task, + hidden_units=[16, 8], + ) + ppnet = create_test_module(ppnet, graph_type) + batch_size = 4 + main_emb = torch.randn(batch_size, main_feature) + uia_emb = torch.randn(batch_size, uia_feature) + + task_outputs = ppnet(main_emb, uia_emb) + + self.assertEqual(len(task_outputs), num_task) + for output in task_outputs: + self.assertEqual( + output.shape, (batch_size, 8) + ) # Output dim from hidden_units + + +class PEPNetTest(unittest.TestCase): + """Test PEPNet model.""" + + @parameterized.expand( + [ + [TestGraphType.NORMAL, False, False], + [TestGraphType.FX_TRACE, False, False], + [TestGraphType.JIT_SCRIPT, False, False], + [TestGraphType.NORMAL, True, False], + [TestGraphType.FX_TRACE, True, False], + [TestGraphType.JIT_SCRIPT, True, False], + [TestGraphType.NORMAL, False, True], + [TestGraphType.FX_TRACE, False, True], + [TestGraphType.JIT_SCRIPT, False, True], + [TestGraphType.NORMAL, True, True], + [TestGraphType.FX_TRACE, True, True], + [TestGraphType.JIT_SCRIPT, True, True], + ] + ) + def test_pepnet(self, graph_type, use_epnet, use_ppnet) -> None: + """Test PEPNet forward pass.""" + feature_cfgs = [ + feature_pb2.FeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="cat_a", embedding_dim=16, num_buckets=100 + ) + ), + feature_pb2.FeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="cat_b", embedding_dim=8, num_buckets=1000 + ) + ), + feature_pb2.FeatureConfig( + raw_feature=feature_pb2.RawFeature(feature_name="int_a") + ), + ] + feature_groups = [ + model_pb2.FeatureGroupConfig( + group_name="main", + feature_names=["cat_a", "cat_b", "int_a"], + group_type=model_pb2.FeatureGroupType.DEEP, + ), + ] + + if use_epnet: + feature_cfgs.append( + feature_pb2.FeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="domain", embedding_dim=16, num_buckets=10 + ) + ) + ) + feature_groups.append( + model_pb2.FeatureGroupConfig( + group_name="domain", + feature_names=["domain"], + group_type=model_pb2.FeatureGroupType.DEEP, + ) + ) + if use_ppnet: + feature_cfgs.append( + feature_pb2.FeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="uia", embedding_dim=16, num_buckets=100 + ) + ) + ) + feature_groups.append( + model_pb2.FeatureGroupConfig( + group_name="uia", + feature_names=["uia"], + group_type=model_pb2.FeatureGroupType.DEEP, + ) + ) + + features = create_features(feature_cfgs) + + pepnet_config = multi_task_rank_pb2.PEPNet( + main_group_name="main", + task_towers=[ + tower_pb2.TaskTower( + tower_name="t1", + label_name="label1", + mlp=module_pb2.MLP(hidden_units=[8, 4]), + losses=[ + loss_pb2.LossConfig( + binary_cross_entropy=loss_pb2.BinaryCrossEntropy() + ) + ], + ), + tower_pb2.TaskTower( + tower_name="t2", + label_name="label2", + mlp=module_pb2.MLP(hidden_units=[8, 4]), + losses=[ + loss_pb2.LossConfig( + binary_cross_entropy=loss_pb2.BinaryCrossEntropy() + ) + ], + ), + ], + ) + if use_epnet: + pepnet_config.domain_group_name = "domain" + pepnet_config.epnet_hidden_unit = 8 + pepnet_config.epnet_gamma = 2.0 + if use_ppnet: + pepnet_config.uia_group_name = "uia" + pepnet_config.ppnet_hidden_units[:] = [16, 8] + pepnet_config.ppnet_activation = "nn.ReLU" + pepnet_config.ppnet_dropout_ratio[:] = [0.1, 0.1] + pepnet_config.ppnet_gamma = 2.0 + + model_config = model_pb2.ModelConfig( + feature_groups=feature_groups, + pepnet=pepnet_config, + ) + + pepnet = PEPNet( + model_config=model_config, + features=features, + labels=["label1", "label2"], + ) + init_parameters(pepnet, device=torch.device("cpu")) + pepnet = create_test_model(pepnet, graph_type) + + sparse_keys = ["cat_a", "cat_b"] + sparse_values = torch.tensor([1, 2, 3, 4, 5, 6, 7]) + sparse_lengths = torch.tensor([1, 2, 1, 3]) + if use_epnet: + sparse_keys.append("domain") + sparse_values = torch.cat([sparse_values, torch.tensor([1, 2])]) + sparse_lengths = torch.cat([sparse_lengths, torch.tensor([1, 1])]) + if use_ppnet: + sparse_keys.append("uia") + sparse_values = torch.cat([sparse_values, torch.tensor([3, 4])]) + sparse_lengths = torch.cat([sparse_lengths, torch.tensor([1, 1])]) + + sparse_feature = KeyedJaggedTensor.from_lengths_sync( + keys=sparse_keys, + values=sparse_values, + lengths=sparse_lengths, + ) + dense_feature = KeyedTensor.from_tensor_list( + keys=["int_a"], tensors=[torch.tensor([[0.2], [0.3]])] + ) + + batch = Batch( + dense_features={BASE_DATA_GROUP: dense_feature}, + sparse_features={BASE_DATA_GROUP: sparse_feature}, + labels={}, + ) + if graph_type == TestGraphType.JIT_SCRIPT: + predictions = pepnet(batch.to_dict()) + else: + predictions = pepnet(batch) + + self.assertEqual(predictions["logits_t1"].size(), (2,)) + self.assertEqual(predictions["probs_t1"].size(), (2,)) + self.assertEqual(predictions["logits_t2"].size(), (2,)) + self.assertEqual(predictions["probs_t2"].size(), (2,)) + + @parameterized.expand( + [[TestGraphType.NORMAL], [TestGraphType.FX_TRACE], [TestGraphType.JIT_SCRIPT]] + ) + def test_pepnet_has_sequences(self, graph_type) -> None: + """Test PEPNet with sequence features.""" + feature_cfgs = [ + feature_pb2.FeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="cat_a", embedding_dim=16, num_buckets=100 + ) + ), + feature_pb2.FeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="cat_b", embedding_dim=8, num_buckets=1000 + ) + ), + feature_pb2.FeatureConfig( + raw_feature=feature_pb2.RawFeature(feature_name="int_a") + ), + feature_pb2.FeatureConfig( + sequence_feature=feature_pb2.SequenceFeature( + sequence_name="click_seq", + features=[ + feature_pb2.SeqFeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="cat_a", + expression="item:cat_a", + embedding_dim=16, + num_buckets=100, + ) + ), + feature_pb2.SeqFeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="cat_b", + expression="item:cat_b", + embedding_dim=8, + num_buckets=1000, + ) + ), + feature_pb2.SeqFeatureConfig( + raw_feature=feature_pb2.RawFeature( + feature_name="int_a", + expression="item:int_a", + ) + ), + ], + ) + ), + feature_pb2.FeatureConfig( + id_feature=feature_pb2.IdFeature( + feature_name="domain", embedding_dim=16, num_buckets=10 + ) + ), + ] + features = create_features(feature_cfgs) + feature_groups = [ + model_pb2.FeatureGroupConfig( + group_name="main", + feature_names=["cat_a", "cat_b", "int_a"], + group_type=model_pb2.FeatureGroupType.DEEP, + sequence_groups=[ + model_pb2.SeqGroupConfig( + group_name="click_seq", + feature_names=[ + "cat_a", + "cat_b", + "int_a", + "click_seq__cat_a", + "click_seq__cat_b", + "click_seq__int_a", + ], + ), + ], + sequence_encoders=[ + seq_encoder_pb2.SeqEncoderConfig( + din_encoder=seq_encoder_pb2.DINEncoder( + input="click_seq", + attn_mlp=module_pb2.MLP(hidden_units=[128, 64]), + ) + ), + ], + ), + model_pb2.FeatureGroupConfig( + group_name="domain", + feature_names=["domain"], + group_type=model_pb2.FeatureGroupType.DEEP, + ), + ] + + pepnet_config = multi_task_rank_pb2.PEPNet( + main_group_name="main", + domain_group_name="domain", + epnet_hidden_unit=8, + task_towers=[ + tower_pb2.TaskTower( + tower_name="t1", + label_name="label1", + mlp=module_pb2.MLP(hidden_units=[8, 4]), + losses=[ + loss_pb2.LossConfig( + binary_cross_entropy=loss_pb2.BinaryCrossEntropy() + ) + ], + ), + tower_pb2.TaskTower( + tower_name="t2", + label_name="label2", + mlp=module_pb2.MLP(hidden_units=[8, 4]), + losses=[ + loss_pb2.LossConfig( + binary_cross_entropy=loss_pb2.BinaryCrossEntropy() + ) + ], + ), + ], + ) + + model_config = model_pb2.ModelConfig( + feature_groups=feature_groups, + pepnet=pepnet_config, + ) + + pepnet = PEPNet( + model_config=model_config, + features=features, + labels=["label1", "label2"], + ) + init_parameters(pepnet, device=torch.device("cpu")) + pepnet = create_test_model(pepnet, graph_type) + + # Batch size is 2 + # cat_a: [1, 1] -> 2 values + # cat_b: [1, 1] -> 2 values + # domain: [1, 1] -> 2 values + # click_seq__cat_a: [2, 3] -> 5 values (sequence feature) + # click_seq__cat_b: [2, 3] -> 5 values (sequence feature) + # Total values: 2+2+2+5+5 = 16 + sparse_feature = KeyedJaggedTensor.from_lengths_sync( + keys=[ + "cat_a", + "cat_b", + "domain", + "click_seq__cat_a", + "click_seq__cat_b", + ], + values=torch.tensor(list(range(16))), + lengths=torch.tensor([1, 1, 1, 1, 1, 1, 2, 3, 2, 3]), + ) + dense_feature = KeyedTensor.from_tensor_list( + keys=["int_a"], tensors=[torch.tensor([[0.2], [0.3]])] + ) + # click_seq__int_a: [2, 3] -> 5 values for 2 samples + sequence_dense_feature = KeyedJaggedTensor.from_lengths_sync( + keys=["click_seq__int_a"], + values=torch.tensor([[x] for x in range(5)], dtype=torch.float32), + lengths=torch.tensor([2, 3]), + ).to_dict() + + batch = Batch( + dense_features={BASE_DATA_GROUP: dense_feature}, + sparse_features={BASE_DATA_GROUP: sparse_feature}, + sequence_dense_features=sequence_dense_feature, + labels={}, + ) + if graph_type == TestGraphType.JIT_SCRIPT: + predictions = pepnet(batch.to_dict()) + else: + predictions = pepnet(batch) + + self.assertEqual(predictions["logits_t1"].size(), (2,)) + self.assertEqual(predictions["probs_t1"].size(), (2,)) + self.assertEqual(predictions["logits_t2"].size(), (2,)) + self.assertEqual(predictions["probs_t2"].size(), (2,)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tzrec/protos/model.proto b/tzrec/protos/model.proto index 59b4aa15..f3c5c2f0 100644 --- a/tzrec/protos/model.proto +++ b/tzrec/protos/model.proto @@ -69,6 +69,7 @@ message ModelConfig { TDM tdm = 400; RocketLaunching rocket_launching = 500; + PEPNet pepnet = 206; } optional uint32 num_class = 2 [default = 1]; diff --git a/tzrec/protos/models/multi_task_rank.proto b/tzrec/protos/models/multi_task_rank.proto index be26e288..7fe53d81 100644 --- a/tzrec/protos/models/multi_task_rank.proto +++ b/tzrec/protos/models/multi_task_rank.proto @@ -70,3 +70,25 @@ message DlrmHSTU { // instead of locally (local rank). optional bool enable_global_average_loss = 5 [default = true]; } + +message PEPNet { + // all feature group name + required string main_group_name = 1; + // domain feature group name + optional string domain_group_name = 2; + // epnet hidden units + optional uint32 epnet_hidden_unit = 3; + // activation function for epnet + optional float epnet_gamma = 4 [default = 2.0]; + // user,item,author feature group name + optional string uia_group_name = 5; + // ppnet hidden units + repeated uint32 ppnet_hidden_units = 6; + // activation function for ppnet + optional string ppnet_activation = 7 [default = 'nn.ReLU']; + // ratio of dropout + repeated float ppnet_dropout_ratio = 8; + optional float ppnet_gamma = 9 [default = 2.0]; + // task tower + repeated TaskTower task_towers = 10; +}