ncf_keras_main.py 16 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
Nimit Nigania's avatar
Nimit Nigania committed
45
from official.utils.flags import core as flags_core
Shining Sun's avatar
Shining Sun committed
46
47
48
49
50


FLAGS = flags.FLAGS


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


62
63
64
65
66
67
68
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
69

70
71
  def call(self, inputs):
    logits, dup_mask = inputs
guptapriya's avatar
guptapriya committed
72
    in_top_k, metric_weights = metric_fn(logits, dup_mask, self.params)
guptapriya's avatar
guptapriya committed
73
    self.add_metric(self.metric(in_top_k, sample_weight=metric_weights))
guptapriya's avatar
guptapriya committed
74
    return logits
75
76


Shining Sun's avatar
Shining Sun committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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()


92
93
94
95
96
97
98
99
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
100
    self.stopped_epoch = 0
101
102
103
104
105
106
107
108
109

  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
110
      print("Epoch %05d: early stopping" % (self.stopped_epoch + 1))
111
112
113
114
115

  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
116
117
118
      logging.warning("Early stopping conditioned on metric `%s` "
                      "which is not available. Available metrics are: %s",
                      self.monitor, ",".join(list(logs.keys())))
119
120
121
    return monitor_value


Shining Sun's avatar
Shining Sun committed
122
123
def _get_keras_model(params):
  """Constructs and returns the model."""
Haoyu Zhang's avatar
Haoyu Zhang committed
124
  batch_size = params["batch_size"]
Shining Sun's avatar
Shining Sun committed
125

126
127
128
129
  # The input layers are of shape (1, batch_size), to match the size of the
  # input data. The first dimension is needed because the input data are
  # required to be batched to use distribution strategies, and in this case, it
  # is designed to be of batch_size 1 for each replica.
Shining Sun's avatar
Shining Sun committed
130
  user_input = tf.keras.layers.Input(
131
      shape=(batch_size,),
132
      batch_size=params["batches_per_step"],
Shining Sun's avatar
Shining Sun committed
133
      name=movielens.USER_COLUMN,
134
      dtype=tf.int32)
Shining Sun's avatar
Shining Sun committed
135
136

  item_input = tf.keras.layers.Input(
137
      shape=(batch_size,),
138
      batch_size=params["batches_per_step"],
Shining Sun's avatar
Shining Sun committed
139
      name=movielens.ITEM_COLUMN,
140
      dtype=tf.int32)
guptapriya's avatar
guptapriya committed
141

142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
  valid_pt_mask_input = tf.keras.layers.Input(
      shape=(batch_size,),
      batch_size=params["batches_per_step"],
      name=rconst.VALID_POINT_MASK,
      dtype=tf.bool)

  dup_mask_input = tf.keras.layers.Input(
      shape=(batch_size,),
      batch_size=params["batches_per_step"],
      name=rconst.DUPLICATE_MASK,
      dtype=tf.int32)

  label_input = tf.keras.layers.Input(
      shape=(batch_size, 1),
      batch_size=params["batches_per_step"],
      name=rconst.TRAIN_LABEL_KEY,
      dtype=tf.bool)
159
160
161

  base_model = neumf_model.construct_model(
      user_input, item_input, params, need_strip=True)
Shining Sun's avatar
Shining Sun committed
162
163
164

  base_model_output = base_model.output

165
166
167
168
  logits = tf.keras.layers.Lambda(
      lambda x: tf.expand_dims(x, 0),
      name="logits")(base_model_output)

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

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

176
177
178
  """CTL does metric calculation as part of eval_step function"""
  if not params["keras_use_ctl"]:
    softmax_logits = MetricLayer(params)([softmax_logits, dup_mask_input])
179

Shining Sun's avatar
Shining Sun committed
180
  keras_model = tf.keras.Model(
guptapriya's avatar
guptapriya committed
181
182
183
184
185
186
      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
187
188
      outputs=softmax_logits)

189
190
191
192
193
194
195
  loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(
      from_logits=True,
      reduction="sum")

  keras_model.add_loss(loss_obj(
      y_true=label_input,
      y_pred=softmax_logits,
guptapriya's avatar
guptapriya committed
196
      sample_weight=valid_pt_mask_input) * 1.0 / batch_size)
197

Shining Sun's avatar
Shining Sun committed
198
199
200
201
202
  keras_model.summary()
  return keras_model


def run_ncf(_):
203
204
  """Run NCF training and eval with Keras."""

205
206
  keras_utils.set_session_config(enable_xla=FLAGS.enable_xla)

guptapriya's avatar
guptapriya committed
207
208
209
  if FLAGS.seed is not None:
    print("Setting tf seed")
    tf.random.set_seed(FLAGS.seed)
210

Shining Sun's avatar
Shining Sun committed
211
212
  # TODO(seemuch): Support different train and eval batch sizes
  if FLAGS.eval_batch_size != FLAGS.batch_size:
213
    logging.warning(
Shining Sun's avatar
Shining Sun committed
214
215
216
217
218
219
        "The Keras implementation of NCF currently does not support batch_size "
        "!= eval_batch_size ({} vs. {}). Overriding eval_batch_size to match "
        "batch_size".format(FLAGS.eval_batch_size, FLAGS.batch_size)
        )
    FLAGS.eval_batch_size = FLAGS.batch_size

Shining Sun's avatar
Shining Sun committed
220
  params = ncf_common.parse_flags(FLAGS)
221
  model_helpers.apply_clean(flags.FLAGS)
Shining Sun's avatar
Shining Sun committed
222

223
224
225
226
227
  strategy = distribution_utils.get_distribution_strategy(
      distribution_strategy=FLAGS.distribution_strategy,
      num_gpus=FLAGS.num_gpus)
  params["distribute_strategy"] = strategy

228
  if not keras_utils.is_v2_0() and strategy is not None:
229
230
    logging.error("NCF Keras only works with distribution strategy in TF 2.0")
    return
guptapriya's avatar
guptapriya committed
231

guptapriya's avatar
guptapriya committed
232
  if (params["keras_use_ctl"] and (
233
      not keras_utils.is_v2_0() or strategy is None)):
234
    logging.error(
guptapriya's avatar
guptapriya committed
235
        "Custom training loop only works with tensorflow 2.0 and dist strat.")
236
237
    return

Shining Sun's avatar
Shining Sun committed
238
  # ncf_common rounds eval_batch_size (this is needed due to a reshape during
239
240
  # eval). This carries over that rounding to batch_size as well. This is the
  # per device batch size
Haoyu Zhang's avatar
Haoyu Zhang committed
241
  params["batch_size"] = params["eval_batch_size"]
242
  batch_size = params["batch_size"]
Shining Sun's avatar
Shining Sun committed
243

244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
  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.
    num_users, num_items, num_train_steps, num_eval_steps, producer = (
        ncf_common.get_inputs(params))
    producer.start()
    per_epoch_callback = IncrementEpochCallback(producer)
    callbacks.append(per_epoch_callback)
  else:
    assert params["eval_dataset_path"] and params["input_meta_data_path"]
259
    with tf.io.gfile.GFile(params["input_meta_data_path"], "rb") as reader:
260
261
262
      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
263
264

  params["num_users"], params["num_items"] = num_users, num_items
265
266
267
268
  (train_input_dataset, eval_input_dataset, num_train_steps, num_eval_steps) = \
      (ncf_input_pipeline.create_ncf_input_data(
          params, producer, input_meta_data))
  steps_per_epoch = None if generate_input_online else num_train_steps
269
270
271

  if FLAGS.early_stopping:
    early_stopping_callback = CustomEarlyStopping(
guptapriya's avatar
guptapriya committed
272
        "val_HR_METRIC", desired_value=FLAGS.hr_threshold)
273
    callbacks.append(early_stopping_callback)
274
275
  with distribution_utils.get_strategy_scope(strategy):
    keras_model = _get_keras_model(params)
276
277
278
279
280
    optimizer = tf.keras.optimizers.Adam(
        learning_rate=params["learning_rate"],
        beta_1=params["beta1"],
        beta_2=params["beta2"],
        epsilon=params["epsilon"])
Nimit Nigania's avatar
Nimit Nigania committed
281
282
283
    if FLAGS.dtype == "fp16":
      optimizer = tf.compat.v1.train.experimental.enable_mixed_precision_graph_rewrite(
        optimizer, loss_scale=flags_core.get_loss_scale(FLAGS, default_for_fp16="dynamic"))
284

Haoyu Zhang's avatar
Haoyu Zhang committed
285
  if params["keras_use_ctl"]:
286
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
287
        reduction="sum",
288
289
290
291
292
293
        from_logits=True)
    train_input_iterator = strategy.make_dataset_iterator(train_input_dataset)
    eval_input_iterator = strategy.make_dataset_iterator(eval_input_dataset)

    def train_step():
      """Called once per step to train the model."""
guptapriya's avatar
guptapriya committed
294
      def step_fn(features):
295
296
        """Computes loss and applied gradient per replica."""
        with tf.GradientTape() as tape:
guptapriya's avatar
guptapriya committed
297
          softmax_logits = keras_model(features)
guptapriya's avatar
guptapriya committed
298
          labels = features[rconst.TRAIN_LABEL_KEY]
299
300
301
          loss = loss_object(labels, softmax_logits,
                             sample_weight=features[rconst.VALID_POINT_MASK])
          loss *= (1.0 / (batch_size*strategy.num_replicas_in_sync))
Nimit Nigania's avatar
Nimit Nigania committed
302
303
          if FLAGS.dtype == "fp16":
            loss = optimizer.get_scaled_loss(loss)
304
305

        grads = tape.gradient(loss, keras_model.trainable_variables)
Nimit Nigania's avatar
Nimit Nigania committed
306
307
        if FLAGS.dtype == "fp16":
          grads = optimizer.get_unscaled_gradients(grads)
308
        # Converting gradients to dense form helps in perf on GPU for NCF
309
310
        grads = neumf_model.sparse_to_dense_grads(
            list(zip(grads, keras_model.trainable_variables)))
311
        optimizer.apply_gradients(grads)
312
313
314
315
316
317
318
319
320
321
        return loss

      per_replica_losses = strategy.experimental_run(step_fn,
                                                     train_input_iterator)
      mean_loss = strategy.reduce(
          tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
      return mean_loss

    def eval_step():
      """Called once per eval step to compute eval metrics."""
guptapriya's avatar
guptapriya committed
322
      def step_fn(features):
323
        """Computes eval metrics per replica."""
guptapriya's avatar
guptapriya committed
324
        softmax_logits = keras_model(features)
guptapriya's avatar
guptapriya committed
325
        in_top_k, metric_weights = metric_fn(
guptapriya's avatar
guptapriya committed
326
            softmax_logits, features[rconst.DUPLICATE_MASK], params)
327
328
329
330
331
332
333
334
335
336
337
338
        hr_sum = tf.reduce_sum(in_top_k*metric_weights)
        hr_count = tf.reduce_sum(metric_weights)
        return hr_sum, hr_count

      per_replica_hr_sum, per_replica_hr_count = (
          strategy.experimental_run(step_fn, eval_input_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

339
340
341
342
    if not FLAGS.run_eagerly:
      train_step = tf.function(train_step)
      eval_step = tf.function(eval_step)

343
344
    time_callback.on_train_begin()
    for epoch in range(FLAGS.train_epochs):
345
346
347
348
349
350
351
352
353
354
355
      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.initialize()

356
357
358
359
360
      train_loss = 0
      for step in range(num_train_steps):
        time_callback.on_batch_begin(step+epoch*num_train_steps)
        train_loss += train_step()
        time_callback.on_batch_end(step+epoch*num_train_steps)
361
      train_loss /= num_train_steps
Haoyu Zhang's avatar
Haoyu Zhang committed
362
      logging.info("Done training epoch %s, epoch loss=%s.",
363
                   epoch+1, train_loss)
364
365
366
367
368
369
370
      eval_input_iterator.initialize()
      hr_sum = 0
      hr_count = 0
      for _ in range(num_eval_steps):
        step_hr_sum, step_hr_count = eval_step()
        hr_sum += step_hr_sum
        hr_count += step_hr_count
Haoyu Zhang's avatar
Haoyu Zhang committed
371
      logging.info("Done eval epoch %s, hr=%s.", epoch+1, hr_sum/hr_count)
372
373
374
375
376
377
378
379
380
381

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

    time_callback.on_train_end()
    eval_results = [None, hr_sum/hr_count]

  else:
    with distribution_utils.get_strategy_scope(strategy):
382
383
384
385
386
387
388
389
390
391
392
      # 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)
393

394
395
396
397
398
399
400
401
      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)
402
403
404
405

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

      eval_results = keras_model.evaluate(
406
          eval_input_dataset, steps=num_eval_steps, verbose=2)
407
408
409
410
411

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

    if history and history.history:
      train_history = history.history
Haoyu Zhang's avatar
Haoyu Zhang committed
412
      train_loss = train_history["loss"][-1]
413

guptapriya's avatar
cleanup  
guptapriya committed
414
  stats = build_stats(train_loss, eval_results, time_callback)
415
416
417
  return stats


418
def build_stats(loss, eval_result, time_callback):
419
420
  """Normalizes and returns dictionary of stats.

Haoyu Zhang's avatar
Haoyu Zhang committed
421
422
423
424
425
426
427
428
  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.
429
430
  """
  stats = {}
431
  if loss:
Haoyu Zhang's avatar
Haoyu Zhang committed
432
    stats["loss"] = loss
433
434

  if eval_result:
Haoyu Zhang's avatar
Haoyu Zhang committed
435
436
    stats["eval_loss"] = eval_result[0]
    stats["eval_hit_rate"] = eval_result[1]
437
438
439

  if time_callback:
    timestamp_log = time_callback.timestamp_log
Haoyu Zhang's avatar
Haoyu Zhang committed
440
441
    stats["step_timestamp_log"] = timestamp_log
    stats["train_finish_time"] = time_callback.train_finish_time
442
    if len(timestamp_log) > 1:
Haoyu Zhang's avatar
Haoyu Zhang committed
443
      stats["avg_exp_per_second"] = (
444
445
446
447
448
          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
449
450
451
452
453
454
455
456
457
458
459
460
461
462


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])
    if FLAGS.tpu:
      raise ValueError("NCF in Keras does not support TPU for now")
    run_ncf(FLAGS)


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