ncf_keras_main.py 16.2 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
25
26
27
28
29
# 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

import os

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

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


FLAGS = flags.FLAGS


guptapriya's avatar
guptapriya committed
48
49
50
51
52
53
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
54
      params["match_mlperf"])
guptapriya's avatar
guptapriya committed
55
56
57
58
  metric_weights = tf.cast(metric_weights, tf.float32)
  return in_top_k, metric_weights


59
60
61
62
63
64
65
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
66

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


guptapriya's avatar
guptapriya committed
74
def _get_train_and_eval_data(producer, params):
Shining Sun's avatar
Shining Sun committed
75
76
  """Returns the datasets for training and evalutating."""

77
78
79
  def preprocess_train_input(features, labels):
    """Pre-process the training data.

80
    This is needed because
81
    - The label needs to be extended to be used in the loss fn
82
83
    - We need the same inputs for training and eval so adding fake inputs
      for DUPLICATE_MASK in training data.
84
85
    """
    labels = tf.expand_dims(labels, -1)
86
87
88
    fake_dup_mask = tf.zeros_like(features[movielens.USER_COLUMN])
    features[rconst.DUPLICATE_MASK] = fake_dup_mask
    features[rconst.TRAIN_LABEL_KEY] = labels
89

90
    if params["distribute_strategy"] or not keras_utils.is_v2_0():
91
92
93
94
      return features
    else:
      # b/134708104
      return (features,)
95

Shining Sun's avatar
Shining Sun committed
96
  train_input_fn = producer.make_input_fn(is_training=True)
97
98
  train_input_dataset = train_input_fn(params).map(
      preprocess_train_input)
99

Shining Sun's avatar
Shining Sun committed
100
  def preprocess_eval_input(features):
101
102
    """Pre-process the eval data.

103
    This is needed because:
104
    - The label needs to be extended to be used in the loss fn
105
106
    - We need the same inputs for training and eval so adding fake inputs
      for VALID_PT_MASK in eval data.
107
    """
108
    labels = tf.cast(tf.zeros_like(features[movielens.USER_COLUMN]), tf.bool)
109
    labels = tf.expand_dims(labels, -1)
110
    fake_valid_pt_mask = tf.cast(
111
        tf.zeros_like(features[movielens.USER_COLUMN]), tf.bool)
112
    features[rconst.VALID_POINT_MASK] = fake_valid_pt_mask
113
    features[rconst.TRAIN_LABEL_KEY] = labels
114

115
    if params["distribute_strategy"] or not keras_utils.is_v2_0():
116
117
118
119
      return features
    else:
      # b/134708104
      return (features,)
Shining Sun's avatar
Shining Sun committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

  eval_input_fn = producer.make_input_fn(is_training=False)
  eval_input_dataset = eval_input_fn(params).map(
      lambda features: preprocess_eval_input(features))

  return train_input_dataset, eval_input_dataset


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()


143
144
145
146
147
148
149
150
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
151
    self.stopped_epoch = 0
152
153
154
155
156
157
158
159
160

  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
161
      print("Epoch %05d: early stopping" % (self.stopped_epoch + 1))
162
163
164
165
166

  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
167
168
169
      logging.warning("Early stopping conditioned on metric `%s` "
                      "which is not available. Available metrics are: %s",
                      self.monitor, ",".join(list(logs.keys())))
170
171
172
    return monitor_value


Shining Sun's avatar
Shining Sun committed
173
174
def _get_keras_model(params):
  """Constructs and returns the model."""
Haoyu Zhang's avatar
Haoyu Zhang committed
175
  batch_size = params["batch_size"]
Shining Sun's avatar
Shining Sun committed
176

177
178
179
180
  # 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
181
  user_input = tf.keras.layers.Input(
182
      shape=(batch_size,),
183
      batch_size=params["batches_per_step"],
Shining Sun's avatar
Shining Sun committed
184
      name=movielens.USER_COLUMN,
185
      dtype=tf.int32)
Shining Sun's avatar
Shining Sun committed
186
187

  item_input = tf.keras.layers.Input(
188
      shape=(batch_size,),
189
      batch_size=params["batches_per_step"],
Shining Sun's avatar
Shining Sun committed
190
      name=movielens.ITEM_COLUMN,
191
      dtype=tf.int32)
guptapriya's avatar
guptapriya committed
192

193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
  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)
210
211
212

  base_model = neumf_model.construct_model(
      user_input, item_input, params, need_strip=True)
Shining Sun's avatar
Shining Sun committed
213
214
215

  base_model_output = base_model.output

216
217
218
219
  logits = tf.keras.layers.Lambda(
      lambda x: tf.expand_dims(x, 0),
      name="logits")(base_model_output)

Shining Sun's avatar
Shining Sun committed
220
  zeros = tf.keras.layers.Lambda(
221
      lambda x: x * 0)(logits)
Shining Sun's avatar
Shining Sun committed
222
223

  softmax_logits = tf.keras.layers.concatenate(
224
      [zeros, logits],
Shining Sun's avatar
Shining Sun committed
225
226
      axis=-1)

227
228
  softmax_logits = MetricLayer(params)([softmax_logits, dup_mask_input])

Shining Sun's avatar
Shining Sun committed
229
  keras_model = tf.keras.Model(
guptapriya's avatar
guptapriya committed
230
231
232
233
234
235
      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
236
237
      outputs=softmax_logits)

238
239
240
241
242
243
244
  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
245
      sample_weight=valid_pt_mask_input) * 1.0 / batch_size)
246

Shining Sun's avatar
Shining Sun committed
247
248
249
250
251
  keras_model.summary()
  return keras_model


def run_ncf(_):
252
253
  """Run NCF training and eval with Keras."""

254
255
  keras_utils.set_session_config(enable_xla=FLAGS.enable_xla)

guptapriya's avatar
guptapriya committed
256
257
258
  if FLAGS.seed is not None:
    print("Setting tf seed")
    tf.random.set_seed(FLAGS.seed)
259

Shining Sun's avatar
Shining Sun committed
260
261
  # TODO(seemuch): Support different train and eval batch sizes
  if FLAGS.eval_batch_size != FLAGS.batch_size:
262
    logging.warning(
Shining Sun's avatar
Shining Sun committed
263
264
265
266
267
268
        "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
269
270
  params = ncf_common.parse_flags(FLAGS)

271
272
273
274
275
  strategy = distribution_utils.get_distribution_strategy(
      distribution_strategy=FLAGS.distribution_strategy,
      num_gpus=FLAGS.num_gpus)
  params["distribute_strategy"] = strategy

guptapriya's avatar
guptapriya committed
276
  if (params["keras_use_ctl"] and (
277
      not keras_utils.is_v2_0() or strategy is None)):
278
    logging.error(
guptapriya's avatar
guptapriya committed
279
        "Custom training loop only works with tensorflow 2.0 and dist strat.")
280
281
    return

Shining Sun's avatar
Shining Sun committed
282
  # ncf_common rounds eval_batch_size (this is needed due to a reshape during
283
284
  # 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
285
  params["batch_size"] = params["eval_batch_size"]
286
  batch_size = params["batch_size"]
Shining Sun's avatar
Shining Sun committed
287

Shining Sun's avatar
Shining Sun committed
288
289
290
291
292
293
294
  num_users, num_items, num_train_steps, num_eval_steps, producer = (
      ncf_common.get_inputs(params))

  params["num_users"], params["num_items"] = num_users, num_items
  producer.start()
  model_helpers.apply_clean(flags.FLAGS)

295
296
297
298
299
300
301
302
303
  batches_per_step = params["batches_per_step"]
  train_input_dataset, eval_input_dataset = _get_train_and_eval_data(producer,
                                                                     params)
  # It is required that for distributed training, the dataset must call
  # batch(). The parameter of batch() here is the number of replicas involed,
  # such that each replica evenly gets a slice of data.
  train_input_dataset = train_input_dataset.batch(batches_per_step)
  eval_input_dataset = eval_input_dataset.batch(batches_per_step)

304
  time_callback = keras_utils.TimeHistory(batch_size, FLAGS.log_steps)
guptapriya's avatar
guptapriya committed
305
306
  per_epoch_callback = IncrementEpochCallback(producer)
  callbacks = [per_epoch_callback, time_callback]
307
308
309

  if FLAGS.early_stopping:
    early_stopping_callback = CustomEarlyStopping(
guptapriya's avatar
guptapriya committed
310
        "val_HR_METRIC", desired_value=FLAGS.hr_threshold)
311
312
    callbacks.append(early_stopping_callback)

313

314
315
  with distribution_utils.get_strategy_scope(strategy):
    keras_model = _get_keras_model(params)
316
317
318
319
320
    optimizer = tf.keras.optimizers.Adam(
        learning_rate=params["learning_rate"],
        beta_1=params["beta1"],
        beta_2=params["beta2"],
        epsilon=params["epsilon"])
321

Haoyu Zhang's avatar
Haoyu Zhang committed
322
  if params["keras_use_ctl"]:
323
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
324
325
326
327
328
329
330
331
        reduction=tf.keras.losses.Reduction.SUM,
        from_logits=True)
    train_input_iterator = strategy.make_dataset_iterator(train_input_dataset)
    eval_input_iterator = strategy.make_dataset_iterator(eval_input_dataset)

    @tf.function
    def train_step():
      """Called once per step to train the model."""
guptapriya's avatar
guptapriya committed
332
      def step_fn(features):
333
334
        """Computes loss and applied gradient per replica."""
        with tf.GradientTape() as tape:
guptapriya's avatar
guptapriya committed
335
          softmax_logits = keras_model(features)
guptapriya's avatar
guptapriya committed
336
          labels = features[rconst.TRAIN_LABEL_KEY]
337
338
339
340
341
          loss = loss_object(labels, softmax_logits,
                             sample_weight=features[rconst.VALID_POINT_MASK])
          loss *= (1.0 / (batch_size*strategy.num_replicas_in_sync))

        grads = tape.gradient(loss, keras_model.trainable_variables)
342
343
344
        # 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)
345
346
347
348
349
350
351
352
353
354
355
        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

    @tf.function
    def eval_step():
      """Called once per eval step to compute eval metrics."""
guptapriya's avatar
guptapriya committed
356
      def step_fn(features):
357
        """Computes eval metrics per replica."""
guptapriya's avatar
guptapriya committed
358
        softmax_logits = keras_model(features)
guptapriya's avatar
guptapriya committed
359
        in_top_k, metric_weights = metric_fn(
guptapriya's avatar
guptapriya committed
360
            softmax_logits, features[rconst.DUPLICATE_MASK], params)
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
        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

    time_callback.on_train_begin()
    for epoch in range(FLAGS.train_epochs):
      per_epoch_callback.on_epoch_begin(epoch)
      train_input_iterator.initialize()
      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)
382
      train_loss /= num_train_steps
Haoyu Zhang's avatar
Haoyu Zhang committed
383
      logging.info("Done training epoch %s, epoch loss=%s.",
384
                   epoch+1, train_loss)
385
386
387
388
389
390
391
      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
392
      logging.info("Done eval epoch %s, hr=%s.", epoch+1, hr_sum/hr_count)
393
394
395
396
397
398
399
400
401
402
403

      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):

404
405
      keras_model.compile(optimizer=optimizer,
                          run_eagerly=FLAGS.run_eagerly)
406
407
408
409
410
411

      history = keras_model.fit(train_input_dataset,
                                epochs=FLAGS.train_epochs,
                                callbacks=callbacks,
                                validation_data=eval_input_dataset,
                                validation_steps=num_eval_steps,
guptapriya's avatar
cleanup  
guptapriya committed
412
                                verbose=2)
413
414
415
416
417
418
419
420
421
422
423
424

      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
Haoyu Zhang's avatar
Haoyu Zhang committed
425
      train_loss = train_history["loss"][-1]
426

guptapriya's avatar
cleanup  
guptapriya committed
427
  stats = build_stats(train_loss, eval_results, time_callback)
428
429
430
  return stats


431
def build_stats(loss, eval_result, time_callback):
432
433
  """Normalizes and returns dictionary of stats.

Haoyu Zhang's avatar
Haoyu Zhang committed
434
435
436
437
438
439
440
441
  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.
442
443
  """
  stats = {}
444
  if loss:
Haoyu Zhang's avatar
Haoyu Zhang committed
445
    stats["loss"] = loss
446
447

  if eval_result:
Haoyu Zhang's avatar
Haoyu Zhang committed
448
449
    stats["eval_loss"] = eval_result[0]
    stats["eval_hit_rate"] = eval_result[1]
450
451
452

  if time_callback:
    timestamp_log = time_callback.timestamp_log
Haoyu Zhang's avatar
Haoyu Zhang committed
453
454
    stats["step_timestamp_log"] = timestamp_log
    stats["train_finish_time"] = time_callback.train_finish_time
455
    if len(timestamp_log) > 1:
Haoyu Zhang's avatar
Haoyu Zhang committed
456
      stats["avg_exp_per_second"] = (
457
458
459
460
461
          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
462
463
464
465
466
467
468
469
470
471
472
473
474
475


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)