From bd3e5796e18f80a420e7923741899192dfe04a30 Mon Sep 17 00:00:00 2001 From: zhouneng2 Date: Thu, 9 Sep 2021 09:28:25 +0800 Subject: [PATCH] improve NCF training speed --- .../official/recommend/ncf/src/dataset.py | 21 +++++++------------ 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/model_zoo/official/recommend/ncf/src/dataset.py b/model_zoo/official/recommend/ncf/src/dataset.py index a7d22ed0b11..f710b7ec16f 100644 --- a/model_zoo/official/recommend/ncf/src/dataset.py +++ b/model_zoo/official/recommend/ncf/src/dataset.py @@ -579,20 +579,13 @@ def create_dataset(test_train=True, data_dir='./dataset/', dataset='ml-1m', trai sampler = RandomSampler(train_pos_users.shape[0], num_neg, batch_size) if rank_id is not None and rank_size is not None: sampler = DistributedSamplerOfTrain(train_pos_users.shape[0], num_neg, batch_size, rank_id, rank_size) - if dataset == 'ml-20m': - ds = GeneratorDataset(dataset, - column_names=[movielens.USER_COLUMN, - movielens.ITEM_COLUMN, - "labels", - rconst.VALID_POINT_MASK], - sampler=sampler, num_parallel_workers=32, python_multiprocessing=False) - else: - ds = GeneratorDataset(dataset, - column_names=[movielens.USER_COLUMN, - movielens.ITEM_COLUMN, - "labels", - rconst.VALID_POINT_MASK], - sampler=sampler) + + ds = GeneratorDataset(dataset, + column_names=[movielens.USER_COLUMN, + movielens.ITEM_COLUMN, + "labels", + rconst.VALID_POINT_MASK], + sampler=sampler) else: eval_batch_size = parse_eval_batch_size(eval_batch_size=eval_batch_size)