MTRec is a simple multi-task recommendation package based Tensorflow2.x. You can install it to invoke some classic models, such as MMoE, and we also provide a test file
to help you better use it.
Let's get started!
MTRec is on PyPI, so you can use pip
to install it.
pip install mtrec==0.0.1
pip install mtrec
The default dependency that MTRec has are Tensorflow2.x.
There are some simple tests to use mtrec package.
First,Dataset.
- MTREC stipulates that features must be sparse discrete features, and continuous features need discrete buckets.
- The dataset is output in the form of a dictionary, for example
{'name1':[n1,n2,...], 'name2':[n1,n2,...]}
Second, Build Model.
from mtrec import MMoE
task_names = ['task1', 'task2']
num_experts = 3
model = MMoE(task_names, num_experts, sparse_feature_columns)
for sparse_feature_columns
,
from mtrec.functions.feature_column import sparseFeature
embed_dim = 4
sparse_feature_columns = [sparseFeature(feat, len(data_df[feat].unique()), embed_dim=embed_dim) for feat in sparse_features]
see the test/utils.py
file for specific details.
Third, Compile, Fit and Predict
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.metrics import AUC
model.compile(loss={'task1': 'binary_crossentropy', 'task2': 'binary_crossentropy'},
optimizer=Adam(learning_rate=learning_rate),
metrics=[AUC()])
model.fit(
train_X,
train_y,
epochs=epochs,
batch_size=batch_size,
)
pred = model.predict(test_X, batch_size=batch_size)
Model | Paper | Published | Author/Group |
---|---|---|---|
MMoE | Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture-of-Experts | KDD, 2018 |
- If you have any suggestions or questions about the project, you can leave a comment on
Issue
or emailzggzy1996@163.com
. - wechat: