ncf_keras_main.py 16.9 KB
Newer Older
Shining Sun's avatar
Shining Sun committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# 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.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

25
import json
Shining Sun's avatar
Shining Sun committed
26
27
28
import os

# pylint: disable=g-bad-import-order
29
from absl import app as absl_app
Shining Sun's avatar
Shining Sun committed
30
from absl import flags
31
from absl import logging
Shining Sun's avatar
Shining Sun committed
32
33
34
35
import tensorflow as tf
# pylint: enable=g-bad-import-order

from official.datasets import movielens
36
from official.recommendation import constants as rconst
Shining Sun's avatar
Shining Sun committed
37
from official.recommendation import ncf_common
38
from official.recommendation import ncf_input_pipeline
Shining Sun's avatar
Shining Sun committed
39
40
41
from official.recommendation import neumf_model
from official.utils.logs import logger
from official.utils.logs import mlperf_helper
42
from official.utils.misc import distribution_utils
43
from official.utils.misc import keras_utils
Shining Sun's avatar
Shining Sun committed
44
from official.utils.misc import model_helpers
45
from official.utils.misc import tpu_lib
Shining Sun's avatar
Shining Sun committed
46
47
48
49

FLAGS = flags.FLAGS


guptapriya's avatar
guptapriya committed
50
51
def metric_fn(logits, dup_mask, params):
  dup_mask = tf.cast(dup_mask, tf.float32)
52
  logits = tf.slice(logits, [0, 1], [-1, -1])
guptapriya's avatar
guptapriya committed
53
54
55
  in_top_k, _, metric_weights, _ = neumf_model.compute_top_k_and_ndcg(
      logits,
      dup_mask,
guptapriya's avatar
cleanup  
guptapriya committed
56
      params["match_mlperf"])
guptapriya's avatar
guptapriya committed
57
58
59
60
  metric_weights = tf.cast(metric_weights, tf.float32)
  return in_top_k, metric_weights


61
62
63
64
65
66
67
class MetricLayer(tf.keras.layers.Layer):
  """Custom layer of metrics for NCF model."""

  def __init__(self, params):
    super(MetricLayer, self).__init__()
    self.params = params
    self.metric = tf.keras.metrics.Mean(name=rconst.HR_METRIC_NAME)
guptapriya's avatar
guptapriya committed
68

69
  def call(self, inputs, training=False):
70
    logits, dup_mask = inputs
71
72
73
74
75
76
77
78

    if not training:
      in_top_k, metric_weights = metric_fn(logits, dup_mask, self.params)
      metric = self.metric(in_top_k, sample_weight=metric_weights)
    else:
      metric = 0.0

    self.add_metric(metric, name="ncf_metric", aggregation="mean")
guptapriya's avatar
guptapriya committed
79
    return logits
80
81


82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class LossLayer(tf.keras.layers.Layer):
  """Pass-through loss layer for NCF model."""

  def __init__(self, loss_normalization_factor):
    super(LossLayer, self).__init__()
    self.loss_normalization_factor = loss_normalization_factor
    self.loss = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True, reduction="sum")

  def call(self, inputs):
    logits, labels, valid_pt_mask_input = inputs
    loss = self.loss(
        y_true=labels, y_pred=logits, sample_weight=valid_pt_mask_input)
    loss = loss * (1.0 / self.loss_normalization_factor)
    self.add_loss(loss)
    return logits


Shining Sun's avatar
Shining Sun committed
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class IncrementEpochCallback(tf.keras.callbacks.Callback):
  """A callback to increase the requested epoch for the data producer.

  The reason why we need this is because we can only buffer a limited amount of
  data. So we keep a moving window to represent the buffer. This is to move the
  one of the window's boundaries for each epoch.
  """

  def __init__(self, producer):
    self._producer = producer

  def on_epoch_begin(self, epoch, logs=None):
    self._producer.increment_request_epoch()


115
116
117
118
119
120
121
122
class CustomEarlyStopping(tf.keras.callbacks.Callback):
  """Stop training has reached a desired hit rate."""

  def __init__(self, monitor, desired_value):
    super(CustomEarlyStopping, self).__init__()

    self.monitor = monitor
    self.desired = desired_value
123
    self.stopped_epoch = 0
124
125
126
127
128
129
130
131
132

  def on_epoch_end(self, epoch, logs=None):
    current = self.get_monitor_value(logs)
    if current and current >= self.desired:
      self.stopped_epoch = epoch
      self.model.stop_training = True

  def on_train_end(self, logs=None):
    if self.stopped_epoch > 0:
Haoyu Zhang's avatar
Haoyu Zhang committed
133
      print("Epoch %05d: early stopping" % (self.stopped_epoch + 1))
134
135
136
137
138

  def get_monitor_value(self, logs):
    logs = logs or {}
    monitor_value = logs.get(self.monitor)
    if monitor_value is None:
Haoyu Zhang's avatar
Haoyu Zhang committed
139
140
141
      logging.warning("Early stopping conditioned on metric `%s` "
                      "which is not available. Available metrics are: %s",
                      self.monitor, ",".join(list(logs.keys())))
142
143
144
    return monitor_value


Shining Sun's avatar
Shining Sun committed
145
146
def _get_keras_model(params):
  """Constructs and returns the model."""
Haoyu Zhang's avatar
Haoyu Zhang committed
147
  batch_size = params["batch_size"]
Shining Sun's avatar
Shining Sun committed
148
149

  user_input = tf.keras.layers.Input(
150
      shape=(1,), name=movielens.USER_COLUMN, dtype=tf.int32)
Shining Sun's avatar
Shining Sun committed
151
152

  item_input = tf.keras.layers.Input(
153
      shape=(1,), name=movielens.ITEM_COLUMN, dtype=tf.int32)
guptapriya's avatar
guptapriya committed
154

155
  valid_pt_mask_input = tf.keras.layers.Input(
156
      shape=(1,), name=rconst.VALID_POINT_MASK, dtype=tf.bool)
157
158

  dup_mask_input = tf.keras.layers.Input(
159
      shape=(1,), name=rconst.DUPLICATE_MASK, dtype=tf.int32)
160
161

  label_input = tf.keras.layers.Input(
162
      shape=(1,), name=rconst.TRAIN_LABEL_KEY, dtype=tf.bool)
Shining Sun's avatar
Shining Sun committed
163

164
  base_model = neumf_model.construct_model(user_input, item_input, params)
Shining Sun's avatar
Shining Sun committed
165

166
  logits = base_model.output
167

Shining Sun's avatar
Shining Sun committed
168
  zeros = tf.keras.layers.Lambda(
169
      lambda x: x * 0)(logits)
Shining Sun's avatar
Shining Sun committed
170
171

  softmax_logits = tf.keras.layers.concatenate(
172
      [zeros, logits],
Shining Sun's avatar
Shining Sun committed
173
174
      axis=-1)

175
176
  # Custom training loop calculates loss and metric as a part of
  # training/evaluation step function.
177
178
  if not params["keras_use_ctl"]:
    softmax_logits = MetricLayer(params)([softmax_logits, dup_mask_input])
179
180
181
182
    # TODO(b/134744680): Use model.add_loss() instead once the API is well
    # supported.
    softmax_logits = LossLayer(batch_size)(
        [softmax_logits, label_input, valid_pt_mask_input])
183

Shining Sun's avatar
Shining Sun committed
184
  keras_model = tf.keras.Model(
guptapriya's avatar
guptapriya committed
185
186
187
188
189
190
      inputs={
          movielens.USER_COLUMN: user_input,
          movielens.ITEM_COLUMN: item_input,
          rconst.VALID_POINT_MASK: valid_pt_mask_input,
          rconst.DUPLICATE_MASK: dup_mask_input,
          rconst.TRAIN_LABEL_KEY: label_input},
Shining Sun's avatar
Shining Sun committed
191
192
193
194
195
196
197
      outputs=softmax_logits)

  keras_model.summary()
  return keras_model


def run_ncf(_):
198
199
  """Run NCF training and eval with Keras."""

200
201
  keras_utils.set_session_config(enable_xla=FLAGS.enable_xla)

guptapriya's avatar
guptapriya committed
202
203
204
  if FLAGS.seed is not None:
    print("Setting tf seed")
    tf.random.set_seed(FLAGS.seed)
205

Shining Sun's avatar
Shining Sun committed
206
  params = ncf_common.parse_flags(FLAGS)
207
  model_helpers.apply_clean(flags.FLAGS)
Shining Sun's avatar
Shining Sun committed
208

209
210
  strategy = distribution_utils.get_distribution_strategy(
      distribution_strategy=FLAGS.distribution_strategy,
211
212
      num_gpus=FLAGS.num_gpus,
      tpu_address=FLAGS.tpu)
213
214
  params["distribute_strategy"] = strategy

215
  if not keras_utils.is_v2_0() and strategy is not None:
216
217
    logging.error("NCF Keras only works with distribution strategy in TF 2.0")
    return
guptapriya's avatar
guptapriya committed
218
  if (params["keras_use_ctl"] and (
219
      not keras_utils.is_v2_0() or strategy is None)):
220
    logging.error(
guptapriya's avatar
guptapriya committed
221
        "Custom training loop only works with tensorflow 2.0 and dist strat.")
222
    return
223
224
225
  if params["use_tpu"] and not params["keras_use_ctl"]:
    logging.error("Custom training loop must be used when using TPUStrategy.")
    return
226

227
  batch_size = params["batch_size"]
228
229
230
231
232
233
234
235
  time_callback = keras_utils.TimeHistory(batch_size, FLAGS.log_steps)
  callbacks = [time_callback]

  producer, input_meta_data = None, None
  generate_input_online = params["train_dataset_path"] is None

  if generate_input_online:
    # Start data producing thread.
236
    num_users, num_items, _, _, producer = ncf_common.get_inputs(params)
237
238
239
240
241
    producer.start()
    per_epoch_callback = IncrementEpochCallback(producer)
    callbacks.append(per_epoch_callback)
  else:
    assert params["eval_dataset_path"] and params["input_meta_data_path"]
242
    with tf.io.gfile.GFile(params["input_meta_data_path"], "rb") as reader:
243
244
245
      input_meta_data = json.loads(reader.read().decode("utf-8"))
      num_users = input_meta_data["num_users"]
      num_items = input_meta_data["num_items"]
Shining Sun's avatar
Shining Sun committed
246
247

  params["num_users"], params["num_items"] = num_users, num_items
248
249
250

  if FLAGS.early_stopping:
    early_stopping_callback = CustomEarlyStopping(
guptapriya's avatar
guptapriya committed
251
        "val_HR_METRIC", desired_value=FLAGS.hr_threshold)
252
    callbacks.append(early_stopping_callback)
253

254
255
256
257
  with tf.device(tpu_lib.get_primary_cpu_task(params["use_tpu"])):
    (train_input_dataset, eval_input_dataset,
     num_train_steps, num_eval_steps) = \
      (ncf_input_pipeline.create_ncf_input_data(
258
          params, producer, input_meta_data, strategy))
259
    steps_per_epoch = None if generate_input_online else num_train_steps
260
261

    with distribution_utils.get_strategy_scope(strategy):
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
      keras_model = _get_keras_model(params)
      optimizer = tf.keras.optimizers.Adam(
          learning_rate=params["learning_rate"],
          beta_1=params["beta1"],
          beta_2=params["beta2"],
          epsilon=params["epsilon"])

      if params["keras_use_ctl"]:
        train_loss, eval_results = run_ncf_custom_training(
            params,
            strategy,
            keras_model,
            optimizer,
            callbacks,
            train_input_dataset,
            eval_input_dataset,
            num_train_steps,
            num_eval_steps,
            generate_input_online=generate_input_online)
281
      else:
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
        # TODO(b/138957587): Remove when force_v2_in_keras_compile is on longer
        # a valid arg for this model. Also remove as a valid flag.
        if FLAGS.force_v2_in_keras_compile is not None:
          keras_model.compile(
              optimizer=optimizer,
              run_eagerly=FLAGS.run_eagerly,
              experimental_run_tf_function=FLAGS.force_v2_in_keras_compile)
        else:
          keras_model.compile(
              optimizer=optimizer, run_eagerly=FLAGS.run_eagerly)

        history = keras_model.fit(
            train_input_dataset,
            epochs=FLAGS.train_epochs,
            steps_per_epoch=steps_per_epoch,
            callbacks=callbacks,
            validation_data=eval_input_dataset,
            validation_steps=num_eval_steps,
            verbose=2)

        logging.info("Training done. Start evaluating")

        eval_results = keras_model.evaluate(
            eval_input_dataset, steps=num_eval_steps, verbose=2)

        logging.info("Keras evaluation is done.")

        if history and history.history:
          train_history = history.history
          train_loss = train_history["loss"][-1]

    stats = build_stats(train_loss, eval_results, time_callback)
    return stats


def run_ncf_custom_training(params,
                            strategy,
                            keras_model,
                            optimizer,
                            callbacks,
                            train_input_dataset,
                            eval_input_dataset,
                            num_train_steps,
                            num_eval_steps,
                            generate_input_online=True):
  """Runs custom training loop.

  Args:
    params: Dictionary containing training parameters.
    strategy: Distribution strategy to be used for distributed training.
    keras_model: Model used for training.
    optimizer: Optimizer used for training.
    callbacks: Callbacks to be invoked between batches/epochs.
    train_input_dataset: tf.data.Dataset used for training.
    eval_input_dataset: tf.data.Dataset used for evaluation.
    num_train_steps: Total number of steps to run for training.
    num_eval_steps: Total number of steps to run for evaluation.
    generate_input_online: Whether input data was generated by data producer.
      When data is generated by data producer, then train dataset must be
      re-initialized after every epoch.

  Returns:
    A tuple of train loss and a list of training and evaluation results.
  """
  loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
      reduction="sum", from_logits=True)
  train_input_iterator = iter(
      strategy.experimental_distribute_dataset(train_input_dataset))
350

351
352
  def train_step(train_iterator):
    """Called once per step to train the model."""
353

354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
    def step_fn(features):
      """Computes loss and applied gradient per replica."""
      with tf.GradientTape() as tape:
        softmax_logits = keras_model(features)
        labels = features[rconst.TRAIN_LABEL_KEY]
        loss = loss_object(
            labels,
            softmax_logits,
            sample_weight=features[rconst.VALID_POINT_MASK])
        loss *= (1.0 / params["batch_size"])

      grads = tape.gradient(loss, keras_model.trainable_variables)
      # Converting gradients to dense form helps in perf on GPU for NCF
      grads = neumf_model.sparse_to_dense_grads(
          list(zip(grads, keras_model.trainable_variables)))
      optimizer.apply_gradients(grads)
      return loss

    per_replica_losses = strategy.experimental_run_v2(
        step_fn, args=(next(train_iterator),))
    mean_loss = strategy.reduce(
        tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
    return mean_loss

  def eval_step(eval_iterator):
    """Called once per eval step to compute eval metrics."""

    def step_fn(features):
      """Computes eval metrics per replica."""
      softmax_logits = keras_model(features)
      in_top_k, metric_weights = metric_fn(softmax_logits,
                                           features[rconst.DUPLICATE_MASK],
                                           params)
      hr_sum = tf.reduce_sum(in_top_k * metric_weights)
      hr_count = tf.reduce_sum(metric_weights)
      return hr_sum, hr_count
390

391
392
393
394
395
396
397
398
    per_replica_hr_sum, per_replica_hr_count = (
        strategy.experimental_run_v2(
            step_fn, args=(next(eval_iterator),)))
    hr_sum = strategy.reduce(
        tf.distribute.ReduceOp.SUM, per_replica_hr_sum, axis=None)
    hr_count = strategy.reduce(
        tf.distribute.ReduceOp.SUM, per_replica_hr_count, axis=None)
    return hr_sum, hr_count
399

400
401
402
  if not FLAGS.run_eagerly:
    train_step = tf.function(train_step)
    eval_step = tf.function(eval_step)
403

404
405
  for callback in callbacks:
    callback.on_train_begin()
406

407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
  train_loss = 0
  for epoch in range(FLAGS.train_epochs):
    for cb in callbacks:
      cb.on_epoch_begin(epoch)

    # As NCF dataset is sampled with randomness, not repeating
    # data elements in each epoch has significant impact on
    # convergence. As so, offline-generated TF record files
    # contains all epoch worth of data. Thus we do not need
    # to initialize dataset when reading from tf record files.
    if generate_input_online:
      train_input_iterator = iter(
          strategy.experimental_distribute_dataset(train_input_dataset))

    train_loss = 0
    for step in range(num_train_steps):
      current_step = step + epoch * num_train_steps
      for c in callbacks:
        c.on_batch_begin(current_step)

      train_loss += train_step(train_input_iterator)

      for c in callbacks:
        c.on_batch_end(current_step)

    train_loss /= num_train_steps
    logging.info("Done training epoch %s, epoch loss=%s.", epoch + 1,
                 train_loss)

    eval_input_iterator = iter(
        strategy.experimental_distribute_dataset(eval_input_dataset))
    hr_sum = 0
    hr_count = 0
    for _ in range(num_eval_steps):
      step_hr_sum, step_hr_count = eval_step(eval_input_iterator)
      hr_sum += step_hr_sum
      hr_count += step_hr_count

    logging.info("Done eval epoch %s, hr=%s.", epoch + 1, hr_sum / hr_count)

    if (FLAGS.early_stopping and
        float(hr_sum / hr_count) > params["hr_threshold"]):
      break

  for c in callbacks:
    c.on_train_end()

  return train_loss, [None, hr_sum / hr_count]
455
456


457
def build_stats(loss, eval_result, time_callback):
458
459
  """Normalizes and returns dictionary of stats.

Haoyu Zhang's avatar
Haoyu Zhang committed
460
461
462
463
464
465
466
467
  Args:
    loss: The final loss at training time.
    eval_result: Output of the eval step. Assumes first value is eval_loss and
      second value is accuracy_top_1.
    time_callback: Time tracking callback likely used during keras.fit.

  Returns:
    Dictionary of normalized results.
468
469
  """
  stats = {}
470
  if loss:
Haoyu Zhang's avatar
Haoyu Zhang committed
471
    stats["loss"] = loss
472
473

  if eval_result:
Haoyu Zhang's avatar
Haoyu Zhang committed
474
475
    stats["eval_loss"] = eval_result[0]
    stats["eval_hit_rate"] = eval_result[1]
476
477
478

  if time_callback:
    timestamp_log = time_callback.timestamp_log
Haoyu Zhang's avatar
Haoyu Zhang committed
479
480
    stats["step_timestamp_log"] = timestamp_log
    stats["train_finish_time"] = time_callback.train_finish_time
481
    if len(timestamp_log) > 1:
Haoyu Zhang's avatar
Haoyu Zhang committed
482
      stats["avg_exp_per_second"] = (
483
484
485
486
487
          time_callback.batch_size * time_callback.log_steps *
          (len(time_callback.timestamp_log)-1) /
          (timestamp_log[-1].timestamp - timestamp_log[0].timestamp))

  return stats
Shining Sun's avatar
Shining Sun committed
488
489
490
491
492
493
494
495
496
497
498
499


def main(_):
  with logger.benchmark_context(FLAGS), \
      mlperf_helper.LOGGER(FLAGS.output_ml_perf_compliance_logging):
    mlperf_helper.set_ncf_root(os.path.split(os.path.abspath(__file__))[0])
    run_ncf(FLAGS)


if __name__ == "__main__":
  ncf_common.define_ncf_flags()
  absl_app.run(main)