ncf_keras_main.py 15.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
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
80
  def preprocess_train_input(features, labels):
    """Pre-process the training data.

    This is needed because:
81
82
83
84
    - Distributed training with keras fit does not support extra inputs. The
      current implementation for fit does not use the VALID_POINT_MASK in the
      input, which makes it extra, so it needs to be removed when using keras
      fit.
85
86
87
    - The label needs to be extended to be used in the loss fn
    """
    labels = tf.expand_dims(labels, -1)
88
89
90
    fake_dup_mask = tf.zeros_like(features[movielens.USER_COLUMN])
    features[rconst.DUPLICATE_MASK] = fake_dup_mask
    features[rconst.TRAIN_LABEL_KEY] = labels
91
92
    return features, labels

Shining Sun's avatar
Shining Sun committed
93
  train_input_fn = producer.make_input_fn(is_training=True)
94
95
  train_input_dataset = train_input_fn(params).map(
      preprocess_train_input)
96

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

    This is needed because:
101
102
103
104
    - Distributed training with keras fit does not support extra inputs. The
      current implementation for fit does not use the DUPLICATE_MASK in the
      input, which makes it extra, so it needs to be removed when using keras
      fit.
105
106
    - The label needs to be extended to be used in the loss fn
    """
107
    labels = tf.cast(tf.zeros_like(features[movielens.USER_COLUMN]), tf.bool)
108
    labels = tf.expand_dims(labels, -1)
109
110
111
112
    fake_valit_pt_mask = tf.cast(
        tf.zeros_like(features[movielens.USER_COLUMN]), tf.bool)
    features[rconst.VALID_POINT_MASK] = fake_valit_pt_mask
    features[rconst.TRAIN_LABEL_KEY] = labels
Shining Sun's avatar
Shining Sun committed
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
    return features, labels

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


137
138
139
140
141
142
143
144
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
145
    self.stopped_epoch = 0
146
147
148
149
150
151
152
153
154

  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
155
      print("Epoch %05d: early stopping" % (self.stopped_epoch + 1))
156
157
158
159
160

  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
161
162
163
      logging.warning("Early stopping conditioned on metric `%s` "
                      "which is not available. Available metrics are: %s",
                      self.monitor, ",".join(list(logs.keys())))
164
165
166
    return monitor_value


Shining Sun's avatar
Shining Sun committed
167
168
def _get_keras_model(params):
  """Constructs and returns the model."""
Haoyu Zhang's avatar
Haoyu Zhang committed
169
  batch_size = params["batch_size"]
Shining Sun's avatar
Shining Sun committed
170

171
172
173
174
  # 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
175
  user_input = tf.keras.layers.Input(
176
      shape=(batch_size,),
177
      batch_size=params["batches_per_step"],
Shining Sun's avatar
Shining Sun committed
178
      name=movielens.USER_COLUMN,
179
      dtype=tf.int32)
Shining Sun's avatar
Shining Sun committed
180
181

  item_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.ITEM_COLUMN,
185
      dtype=tf.int32)
guptapriya's avatar
guptapriya committed
186

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
  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)
204
205
206

  base_model = neumf_model.construct_model(
      user_input, item_input, params, need_strip=True)
Shining Sun's avatar
Shining Sun committed
207
208
209

  base_model_output = base_model.output

210
211
212
213
  logits = tf.keras.layers.Lambda(
      lambda x: tf.expand_dims(x, 0),
      name="logits")(base_model_output)

Shining Sun's avatar
Shining Sun committed
214
  zeros = tf.keras.layers.Lambda(
215
      lambda x: x * 0)(logits)
Shining Sun's avatar
Shining Sun committed
216
217

  softmax_logits = tf.keras.layers.concatenate(
218
      [zeros, logits],
Shining Sun's avatar
Shining Sun committed
219
220
      axis=-1)

221
222
  softmax_logits = MetricLayer(params)([softmax_logits, dup_mask_input])

Shining Sun's avatar
Shining Sun committed
223
  keras_model = tf.keras.Model(
guptapriya's avatar
guptapriya committed
224
225
226
227
228
229
      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
230
231
      outputs=softmax_logits)

232
233
234
235
  loss_obj = tf.keras.losses.SparseCategoricalCrossentropy(
      from_logits=True,
      reduction="sum")

guptapriya's avatar
guptapriya committed
236
237
  loss_scale_factor = (batch_size *
                       tf.distribute.get_strategy().num_replicas_in_sync)
238
239
240
  keras_model.add_loss(loss_obj(
      y_true=label_input,
      y_pred=softmax_logits,
guptapriya's avatar
guptapriya committed
241
      sample_weight=valid_pt_mask_input) * 1.0 / loss_scale_factor)
242

Shining Sun's avatar
Shining Sun committed
243
244
245
246
247
248
  keras_model.summary()
  return keras_model


def run_ncf(_):
  """Run NCF training and eval with Keras."""
Shining Sun's avatar
Shining Sun committed
249
250
  # TODO(seemuch): Support different train and eval batch sizes
  if FLAGS.eval_batch_size != FLAGS.batch_size:
251
    logging.warning(
Shining Sun's avatar
Shining Sun committed
252
253
254
255
256
257
        "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
258
259
  params = ncf_common.parse_flags(FLAGS)

Haoyu Zhang's avatar
Haoyu Zhang committed
260
  if params["keras_use_ctl"] and int(tf.__version__.split(".")[0]) == 1:
261
262
263
264
    logging.error(
        "Custom training loop only works with tensorflow 2.0 and above.")
    return

Shining Sun's avatar
Shining Sun committed
265
  # ncf_common rounds eval_batch_size (this is needed due to a reshape during
266
267
  # 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
268
  params["batch_size"] = params["eval_batch_size"]
269
  batch_size = params["batch_size"]
Shining Sun's avatar
Shining Sun committed
270

Shining Sun's avatar
Shining Sun committed
271
272
273
274
275
276
277
  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)

278
279
280
281
282
283
284
285
286
  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)

287
  time_callback = keras_utils.TimeHistory(batch_size, FLAGS.log_steps)
guptapriya's avatar
guptapriya committed
288
289
  per_epoch_callback = IncrementEpochCallback(producer)
  callbacks = [per_epoch_callback, time_callback]
290
291
292
293
294
295

  if FLAGS.early_stopping:
    early_stopping_callback = CustomEarlyStopping(
        "val_metric_fn", desired_value=FLAGS.hr_threshold)
    callbacks.append(early_stopping_callback)

296
297
298
  strategy = ncf_common.get_distribution_strategy(params)
  with distribution_utils.get_strategy_scope(strategy):
    keras_model = _get_keras_model(params)
299
300
301
302
303
    optimizer = tf.keras.optimizers.Adam(
        learning_rate=params["learning_rate"],
        beta_1=params["beta1"],
        beta_2=params["beta2"],
        epsilon=params["epsilon"])
304

Haoyu Zhang's avatar
Haoyu Zhang committed
305
  if params["keras_use_ctl"]:
306
307
308
309
310
311
312
313
314
315
316
317
318
    loss_object = tf.losses.SparseCategoricalCrossentropy(
        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."""
      def step_fn(inputs):
        """Computes loss and applied gradient per replica."""
        features, labels = inputs
        with tf.GradientTape() as tape:
guptapriya's avatar
guptapriya committed
319
          softmax_logits = keras_model(features)
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
          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)
        optimizer.apply_gradients(list(zip(grads,
                                           keras_model.trainable_variables)))
        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."""
      def step_fn(inputs):
        """Computes eval metrics per replica."""
        features, _ = inputs
guptapriya's avatar
guptapriya committed
341
        softmax_logits = keras_model(features)
guptapriya's avatar
guptapriya committed
342
        in_top_k, metric_weights = metric_fn(
guptapriya's avatar
guptapriya committed
343
            softmax_logits, features[rconst.DUPLICATE_MASK], params)
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
        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)
365
      train_loss /= num_train_steps
Haoyu Zhang's avatar
Haoyu Zhang committed
366
      logging.info("Done training epoch %s, epoch loss=%s.",
367
                   epoch+1, train_loss)
368
369
370
371
372
373
374
      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
375
      logging.info("Done eval epoch %s, hr=%s.", epoch+1, hr_sum/hr_count)
376
377
378
379
380
381
382
383
384
385
386

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

guptapriya's avatar
guptapriya committed
387
      keras_model.compile(optimizer=optimizer)
388
389
390
391
392
393

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

      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
407
      train_loss = train_history["loss"][-1]
408

guptapriya's avatar
cleanup  
guptapriya committed
409
  stats = build_stats(train_loss, eval_results, time_callback)
410
411
412
  return stats


413
def build_stats(loss, eval_result, time_callback):
414
415
  """Normalizes and returns dictionary of stats.

Haoyu Zhang's avatar
Haoyu Zhang committed
416
417
418
419
420
421
422
423
  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.
424
425
  """
  stats = {}
426
  if loss:
Haoyu Zhang's avatar
Haoyu Zhang committed
427
    stats["loss"] = loss
428
429

  if eval_result:
Haoyu Zhang's avatar
Haoyu Zhang committed
430
431
    stats["eval_loss"] = eval_result[0]
    stats["eval_hit_rate"] = eval_result[1]
432
433
434

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


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)