ncf_main.py 20 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""NCF framework to train and evaluate the NeuMF model.

The NeuMF model assembles both MF and MLP models under the NCF framework. Check
`neumf_model.py` for more details about the models.
"""
20

21
22
23
24
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

25
import contextlib
26
27
import heapq
import math
28
import multiprocessing
29
import os
30
31
import signal
import typing
32

33
# pylint: disable=g-bad-import-order
34
import numpy as np
35
36
from absl import app as absl_app
from absl import flags
37
import tensorflow as tf
38
# pylint: enable=g-bad-import-order
39

Reed's avatar
Reed committed
40
from tensorflow.contrib.compiler import xla
41
from official.datasets import movielens
42
43
from official.recommendation import constants as rconst
from official.recommendation import data_preprocessing
44
from official.recommendation import model_runner
45
from official.recommendation import neumf_model
46
47
48
from official.utils.flags import core as flags_core
from official.utils.logs import hooks_helper
from official.utils.logs import logger
49
from official.utils.logs import mlperf_helper
50
from official.utils.misc import distribution_utils
51
from official.utils.misc import model_helpers
52
53


Reed's avatar
Reed committed
54
55
56
FLAGS = flags.FLAGS


57
def construct_estimator(num_gpus, model_dir, iterations, params, batch_size,
58
59
                        eval_batch_size):
  """Construct either an Estimator or TPUEstimator for NCF.
60
61

  Args:
62
63
    num_gpus: The number of gpus (Used to select distribution strategy)
    model_dir: The model directory for the estimator
64
    iterations:  Estimator iterations
65
66
67
    params: The params dict for the estimator
    batch_size: The mini-batch size for training.
    eval_batch_size: The batch size used during evaluation.
68
69

  Returns:
70
    An Estimator or TPUEstimator.
71
72
  """

73
74
75
76
77
78
  if params["use_tpu"]:
    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(
        tpu=params["tpu"],
        zone=params["tpu_zone"],
        project=params["tpu_gcp_project"],
    )
79
80
    tf.logging.info("Issuing reset command to TPU to ensure a clean state.")
    tf.Session.reset(tpu_cluster_resolver.get_master())
81
82

    tpu_config = tf.contrib.tpu.TPUConfig(
83
        iterations_per_loop=iterations,
84
85
86
87
88
        num_shards=8)

    run_config = tf.contrib.tpu.RunConfig(
        cluster=tpu_cluster_resolver,
        model_dir=model_dir,
89
        save_checkpoints_secs=600,
90
91
92
93
94
95
96
97
98
99
        session_config=tf.ConfigProto(
            allow_soft_placement=True, log_device_placement=False),
        tpu_config=tpu_config)

    tpu_params = {k: v for k, v in params.items() if k != "batch_size"}

    train_estimator = tf.contrib.tpu.TPUEstimator(
        model_fn=neumf_model.neumf_model_fn,
        use_tpu=True,
        train_batch_size=batch_size,
100
        eval_batch_size=eval_batch_size,
101
102
103
104
105
        params=tpu_params,
        config=run_config)

    eval_estimator = tf.contrib.tpu.TPUEstimator(
        model_fn=neumf_model.neumf_model_fn,
106
        use_tpu=True,
107
        train_batch_size=1,
108
        eval_batch_size=eval_batch_size,
109
110
111
112
113
114
        params=tpu_params,
        config=run_config)

    return train_estimator, eval_estimator

  distribution = distribution_utils.get_distribution_strategy(num_gpus=num_gpus)
115
116
  run_config = tf.estimator.RunConfig(train_distribute=distribution,
                                      eval_distribute=distribution)
117
  params["eval_batch_size"] = eval_batch_size
Reed's avatar
Reed committed
118
119
120
121
122
123
  model_fn = neumf_model.neumf_model_fn
  if params["use_xla_for_gpu"]:
    tf.logging.info("Using XLA for GPU for training and evaluation.")
    model_fn = xla.estimator_model_fn(model_fn)
  estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir=model_dir,
                                     config=run_config, params=params)
124
  return estimator, estimator
125
126
127


def main(_):
Reed's avatar
Reed committed
128
129
  with logger.benchmark_context(FLAGS), \
       mlperf_helper.LOGGER(FLAGS.output_ml_perf_compliance_logging):
130
    mlperf_helper.set_ncf_root(os.path.split(os.path.abspath(__file__))[0])
131
    run_ncf(FLAGS)
132
    mlperf_helper.stitch_ncf()
133
134
135
136


def run_ncf(_):
  """Run NCF training and eval loop."""
137
  if FLAGS.download_if_missing and not FLAGS.use_synthetic_data:
138
    movielens.download(FLAGS.dataset, FLAGS.data_dir)
139

140
141
142
  if FLAGS.seed is not None:
    np.random.seed(FLAGS.seed)

143
144
145
  num_gpus = flags_core.get_num_gpus(FLAGS)
  batch_size = distribution_utils.per_device_batch_size(
      int(FLAGS.batch_size), num_gpus)
146
147

  eval_per_user = rconst.NUM_EVAL_NEGATIVES + 1
Taylor Robie's avatar
Taylor Robie committed
148
149
  eval_batch_size = int(FLAGS.eval_batch_size or
                        max([FLAGS.batch_size, eval_per_user]))
150
151
152
153
154
155
  if eval_batch_size % eval_per_user:
    eval_batch_size = eval_batch_size // eval_per_user * eval_per_user
    tf.logging.warning(
        "eval examples per user does not evenly divide eval_batch_size. "
        "Overriding to {}".format(eval_batch_size))

156
157
158
159
160
  if FLAGS.use_synthetic_data:
    ncf_dataset = None
    cleanup_fn = lambda: None
    num_users, num_items = data_preprocessing.DATASET_TO_NUM_USERS_AND_ITEMS[
        FLAGS.dataset]
161
162
    num_train_steps = data_preprocessing.SYNTHETIC_BATCHES_PER_EPOCH
    num_eval_steps = data_preprocessing.SYNTHETIC_BATCHES_PER_EPOCH
163
164
165
166
167
168
169
170
  else:
    ncf_dataset, cleanup_fn = data_preprocessing.instantiate_pipeline(
        dataset=FLAGS.dataset, data_dir=FLAGS.data_dir,
        batch_size=batch_size,
        eval_batch_size=eval_batch_size,
        num_neg=FLAGS.num_neg,
        epochs_per_cycle=FLAGS.epochs_between_evals,
        match_mlperf=FLAGS.ml_perf,
shizhiw's avatar
shizhiw committed
171
        deterministic=FLAGS.seed is not None,
172
173
        use_subprocess=FLAGS.use_subprocess,
        cache_id=FLAGS.cache_id)
174
175
    num_users = ncf_dataset.num_users
    num_items = ncf_dataset.num_items
176
177
178
179
180
    num_train_steps = int(np.ceil(
        FLAGS.epochs_between_evals * ncf_dataset.num_train_positives *
        (1 + FLAGS.num_neg) / FLAGS.batch_size))
    num_eval_steps = int(np.ceil((1 + rconst.NUM_EVAL_NEGATIVES) *
                                 ncf_dataset.num_users / eval_batch_size))
181
182

  model_helpers.apply_clean(flags.FLAGS)
183

184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
  params = {
      "use_seed": FLAGS.seed is not None,
      "hash_pipeline": FLAGS.hash_pipeline,
      "batch_size": batch_size,
      "eval_batch_size": eval_batch_size,
      "learning_rate": FLAGS.learning_rate,
      "num_users": num_users,
      "num_items": num_items,
      "mf_dim": FLAGS.num_factors,
      "model_layers": [int(layer) for layer in FLAGS.layers],
      "mf_regularization": FLAGS.mf_regularization,
      "mlp_reg_layers": [float(reg) for reg in FLAGS.mlp_regularization],
      "num_neg": FLAGS.num_neg,
      "use_tpu": FLAGS.tpu is not None,
      "tpu": FLAGS.tpu,
      "tpu_zone": FLAGS.tpu_zone,
      "tpu_gcp_project": FLAGS.tpu_gcp_project,
      "beta1": FLAGS.beta1,
      "beta2": FLAGS.beta2,
      "epsilon": FLAGS.epsilon,
      "match_mlperf": FLAGS.ml_perf,
      "use_xla_for_gpu": FLAGS.use_xla_for_gpu,
      "use_estimator": FLAGS.use_estimator,
  }
  if FLAGS.use_estimator:
    train_estimator, eval_estimator = construct_estimator(
210
211
        num_gpus=num_gpus, model_dir=FLAGS.model_dir,
        iterations=num_train_steps, params=params,
212
213
        batch_size=flags.FLAGS.batch_size, eval_batch_size=eval_batch_size)
  else:
Reed's avatar
Reed committed
214
215
    runner = model_runner.NcfModelRunner(ncf_dataset, params, num_train_steps,
                                         num_eval_steps, FLAGS.use_while_loop)
216

217
218
219
  # Create hooks that log information about the training and metric values
  train_hooks = hooks_helper.get_train_hooks(
      FLAGS.hooks,
220
      model_dir=FLAGS.model_dir,
221
222
      batch_size=FLAGS.batch_size,  # for ExamplesPerSecondHook
      tensors_to_log={"cross_entropy": "cross_entropy"}
223
224
225
  )
  run_params = {
      "batch_size": FLAGS.batch_size,
226
      "eval_batch_size": eval_batch_size,
227
228
229
230
      "number_factors": FLAGS.num_factors,
      "hr_threshold": FLAGS.hr_threshold,
      "train_epochs": FLAGS.train_epochs,
  }
231
  benchmark_logger = logger.get_benchmark_logger()
232
233
234
  benchmark_logger.log_run_info(
      model_name="recommendation",
      dataset_name=FLAGS.dataset,
235
236
      run_params=run_params,
      test_id=FLAGS.benchmark_test_id)
237
238


239
  eval_input_fn = None
240
  total_training_cycle = FLAGS.train_epochs // FLAGS.epochs_between_evals
241
242
  target_reached = False
  mlperf_helper.ncf_print(key=mlperf_helper.TAGS.TRAIN_LOOP)
243
  for cycle_index in range(total_training_cycle):
244
    assert FLAGS.epochs_between_evals == 1 or not mlperf_helper.LOGGER.enabled
245
    tf.logging.info("Starting a training cycle: {}/{}".format(
246
        cycle_index + 1, total_training_cycle))
247

248
249
250
    mlperf_helper.ncf_print(key=mlperf_helper.TAGS.TRAIN_EPOCH,
                            value=cycle_index)

251
    # Train the model
252
253
254
255
256
257
    if FLAGS.use_estimator:
      train_input_fn, train_record_dir, batch_count = \
        data_preprocessing.make_input_fn(
            ncf_dataset=ncf_dataset, is_training=True)

      if batch_count != num_train_steps:
258
259
        raise ValueError(
            "Step counts do not match. ({} vs. {}) The async process is "
260
261
262
263
264
265
266
267
            "producing incorrect shards.".format(batch_count, num_train_steps))

      train_estimator.train(input_fn=train_input_fn, hooks=train_hooks,
                            steps=num_train_steps)
      if train_record_dir:
        tf.gfile.DeleteRecursively(train_record_dir)

      tf.logging.info("Beginning evaluation.")
268
269
      if eval_input_fn is None:
        eval_input_fn, _, eval_batch_count = data_preprocessing.make_input_fn(
270
271
272
273
274
275
276
277
278
279
            ncf_dataset=ncf_dataset, is_training=False)

        if eval_batch_count != num_eval_steps:
          raise ValueError(
              "Step counts do not match. ({} vs. {}) The async process is "
              "producing incorrect shards.".format(
                  eval_batch_count, num_eval_steps))

      mlperf_helper.ncf_print(key=mlperf_helper.TAGS.EVAL_START,
                              value=cycle_index)
280
      eval_results = eval_estimator.evaluate(eval_input_fn,
281
282
283
                                             steps=num_eval_steps)
      tf.logging.info("Evaluation complete.")
    else:
Reed's avatar
Reed committed
284
      runner.train()
285
286
287
      tf.logging.info("Beginning evaluation.")
      mlperf_helper.ncf_print(key=mlperf_helper.TAGS.EVAL_START,
                              value=cycle_index)
Reed's avatar
Reed committed
288
      eval_results = runner.eval()
289
      tf.logging.info("Evaluation complete.")
290
291
    hr = float(eval_results[rconst.HR_KEY])
    ndcg = float(eval_results[rconst.NDCG_KEY])
292

293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
    mlperf_helper.ncf_print(
        key=mlperf_helper.TAGS.EVAL_TARGET,
        value={"epoch": cycle_index, "value": FLAGS.hr_threshold})
    mlperf_helper.ncf_print(key=mlperf_helper.TAGS.EVAL_ACCURACY,
                            value={"epoch": cycle_index, "value": hr})
    mlperf_helper.ncf_print(
        key=mlperf_helper.TAGS.EVAL_HP_NUM_NEG,
        value={"epoch": cycle_index, "value": rconst.NUM_EVAL_NEGATIVES})

    # Logged by the async process during record creation.
    mlperf_helper.ncf_print(key=mlperf_helper.TAGS.EVAL_HP_NUM_USERS,
                            deferred=True)

    mlperf_helper.ncf_print(key=mlperf_helper.TAGS.EVAL_STOP, value=cycle_index)

308
309
310
    # Benchmark the evaluation results
    benchmark_logger.log_evaluation_result(eval_results)
    # Log the HR and NDCG results.
311
    tf.logging.info(
312
313
314
315
316
        "Iteration {}: HR = {:.4f}, NDCG = {:.4f}".format(
            cycle_index + 1, hr, ndcg))

    # If some evaluation threshold is met
    if model_helpers.past_stop_threshold(FLAGS.hr_threshold, hr):
317
      target_reached = True
318
319
      break

320
321
  mlperf_helper.ncf_print(key=mlperf_helper.TAGS.RUN_STOP,
                          value={"success": target_reached})
322
323
  cleanup_fn()  # Cleanup data construction artifacts and subprocess.

324
325
326
  # Clear the session explicitly to avoid session delete error
  tf.keras.backend.clear_session()

327
328
  mlperf_helper.ncf_print(key=mlperf_helper.TAGS.RUN_FINAL)

329
330
331
332
333
334
335
336
337

def define_ncf_flags():
  """Add flags for running ncf_main."""
  # Add common flags
  flags_core.define_base(export_dir=False)
  flags_core.define_performance(
      num_parallel_calls=False,
      inter_op=False,
      intra_op=False,
338
      synthetic_data=True,
339
      max_train_steps=False,
340
341
      dtype=False,
      all_reduce_alg=False
342
  )
343
  flags_core.define_device(tpu=True)
344
345
346
347
348
349
350
351
352
  flags_core.define_benchmark()

  flags.adopt_module_key_flags(flags_core)

  flags_core.set_defaults(
      model_dir="/tmp/ncf/",
      data_dir="/tmp/movielens-data/",
      train_epochs=2,
      batch_size=256,
353
354
355
      hooks="ProfilerHook",
      tpu=None
  )
356
357
358
359
360
361
362
363

  # Add ncf-specific flags
  flags.DEFINE_enum(
      name="dataset", default="ml-1m",
      enum_values=["ml-1m", "ml-20m"], case_sensitive=False,
      help=flags_core.help_wrap(
          "Dataset to be trained and evaluated."))

364
365
366
367
  flags.DEFINE_boolean(
      name="download_if_missing", default=True, help=flags_core.help_wrap(
          "Download data to data_dir if it is not already present."))

368
369
370
371
372
373
374
  flags.DEFINE_string(
      name="eval_batch_size", default=None, help=flags_core.help_wrap(
          "The batch size used for evaluation. This should generally be larger"
          "than the training batch size as the lack of back propagation during"
          "evaluation can allow for larger batch sizes to fit in memory. If not"
          "specified, the training batch size (--batch_size) will be used."))

375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
  flags.DEFINE_integer(
      name="num_factors", default=8,
      help=flags_core.help_wrap("The Embedding size of MF model."))

  # Set the default as a list of strings to be consistent with input arguments
  flags.DEFINE_list(
      name="layers", default=["64", "32", "16", "8"],
      help=flags_core.help_wrap(
          "The sizes of hidden layers for MLP. Example "
          "to specify different sizes of MLP layers: --layers=32,16,8,4"))

  flags.DEFINE_float(
      name="mf_regularization", default=0.,
      help=flags_core.help_wrap(
          "The regularization factor for MF embeddings. The factor is used by "
          "regularizer which allows to apply penalties on layer parameters or "
          "layer activity during optimization."))

  flags.DEFINE_list(
      name="mlp_regularization", default=["0.", "0.", "0.", "0."],
      help=flags_core.help_wrap(
          "The regularization factor for each MLP layer. See mf_regularization "
          "help for more info about regularization factor."))

  flags.DEFINE_integer(
      name="num_neg", default=4,
      help=flags_core.help_wrap(
          "The Number of negative instances to pair with a positive instance."))

  flags.DEFINE_float(
      name="learning_rate", default=0.001,
      help=flags_core.help_wrap("The learning rate."))

408
409
410
411
412
413
414
415
416
417
418
419
420
  flags.DEFINE_float(
      name="beta1", default=0.9,
      help=flags_core.help_wrap("beta1 hyperparameter for the Adam optimizer."))

  flags.DEFINE_float(
      name="beta2", default=0.999,
      help=flags_core.help_wrap("beta2 hyperparameter for the Adam optimizer."))

  flags.DEFINE_float(
      name="epsilon", default=1e-8,
      help=flags_core.help_wrap("epsilon hyperparameter for the Adam "
                                "optimizer."))

421
422
423
424
425
426
427
428
  flags.DEFINE_float(
      name="hr_threshold", default=None,
      help=flags_core.help_wrap(
          "If passed, training will stop when the evaluation metric HR is "
          "greater than or equal to hr_threshold. For dataset ml-1m, the "
          "desired hr_threshold is 0.68 which is the result from the paper; "
          "For dataset ml-20m, the threshold can be set as 0.95 which is "
          "achieved by MLPerf implementation."))
429

430
  flags.DEFINE_bool(
431
      name="ml_perf", default=False,
432
433
434
435
436
437
438
439
440
441
442
443
444
      help=flags_core.help_wrap(
          "If set, changes the behavior of the model slightly to match the "
          "MLPerf reference implementations here: \n"
          "https://github.com/mlperf/reference/tree/master/recommendation/"
          "pytorch\n"
          "The two changes are:\n"
          "1. When computing the HR and NDCG during evaluation, remove "
          "duplicate user-item pairs before the computation. This results in "
          "better HRs and NDCGs.\n"
          "2. Use a different soring algorithm when sorting the input data, "
          "which performs better due to the fact the sorting algorithms are "
          "not stable."))

Reed's avatar
Reed committed
445
446
447
448
449
450
451
452
453
454
455
456
  flags.DEFINE_bool(
      name="output_ml_perf_compliance_logging", default=False,
      help=flags_core.help_wrap(
          "If set, output the MLPerf compliance logging. This is only useful "
          "if one is running the model for MLPerf. See "
          "https://github.com/mlperf/policies/blob/master/training_rules.adoc"
          "#submission-compliance-logs for details. This uses sudo and so may "
          "ask for your password, as root access is needed to clear the system "
          "caches, which is required for MLPerf compliance."
      )
  )

457
458
459
460
461
462
463
464
465
466
467
  flags.DEFINE_integer(
      name="seed", default=None, help=flags_core.help_wrap(
          "This value will be used to seed both NumPy and TensorFlow."))

  flags.DEFINE_bool(
      name="hash_pipeline", default=False, help=flags_core.help_wrap(
          "This flag will perform a separate run of the pipeline and hash "
          "batches as they are produced. \nNOTE: this will significantly slow "
          "training. However it is useful to confirm that a random seed is "
          "does indeed make the data pipeline deterministic."))

468
469
470
  @flags.validator("eval_batch_size", "eval_batch_size must be at least {}"
                   .format(rconst.NUM_EVAL_NEGATIVES + 1))
  def eval_size_check(eval_batch_size):
Taylor Robie's avatar
Taylor Robie committed
471
472
    return (eval_batch_size is None or
            int(eval_batch_size) > rconst.NUM_EVAL_NEGATIVES)
473

474
475
476
477
478
479
  flags.DEFINE_bool(
      name="use_subprocess", default=True, help=flags_core.help_wrap(
          "By default, ncf_main.py starts async data generation process as a "
          "subprocess. If set to False, ncf_main.py will assume the async data "
          "generation process has already been started by the user."))

480
481
482
483
484
485
  flags.DEFINE_integer(name="cache_id", default=None, help=flags_core.help_wrap(
      "Use a specified cache_id rather than using a timestamp. This is only "
      "needed to synchronize across multiple workers. Generally this flag will "
      "not need to be set."
  ))

Reed's avatar
Reed committed
486
487
488
489
490
  flags.DEFINE_bool(
      name="use_xla_for_gpu", default=False, help=flags_core.help_wrap(
          "If True, use XLA for the model function. Only works when using a "
          "GPU. On TPUs, XLA is always used"))

491
492
493
494
  xla_message = "--use_xla_for_gpu is incompatible with --tpu"
  @flags.multi_flags_validator(["use_xla_for_gpu", "tpu"], message=xla_message)
  def xla_validator(flag_dict):
    return not flag_dict["use_xla_for_gpu"] or not flag_dict["tpu"]
Reed's avatar
Reed committed
495

496
497
498
499
500
501
502
503
504
  flags.DEFINE_bool(
      name="use_estimator", default=True, help=flags_core.help_wrap(
          "If True, use Estimator to train. Setting to False is slightly "
          "faster, but when False, the following are currently unsupported:\n"
          "  * Using TPUs\n"
          "  * Using more than 1 GPU\n"
          "  * Reloading from checkpoints\n"
          "  * Any hooks specified with --hooks\n"))

Reed's avatar
Reed committed
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
  flags.DEFINE_bool(
      name="use_while_loop", default=None, help=flags_core.help_wrap(
          "If set, run an entire epoch in a session.run() call using a "
          "TensorFlow while loop. This can improve performance, but will not "
          "print out losses throughout the epoch. Requires "
          "--use_estimator=false"
      ))

  xla_message = "--use_while_loop requires --use_estimator=false"
  @flags.multi_flags_validator(["use_while_loop", "use_estimator"],
                               message=xla_message)
  def while_loop_validator(flag_dict):
    return (not flag_dict["use_while_loop"] or
            not flag_dict["use_estimator"])

520
521
522

if __name__ == "__main__":
  tf.logging.set_verbosity(tf.logging.INFO)
523
524
  define_ncf_flags()
  absl_app.run(main)