ncf_common.py 12 KB
Newer Older
Frederick Liu's avatar
Frederick Liu committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Frederick Liu's avatar
Frederick Liu committed
14

Hongkun Yu's avatar
Hongkun Yu committed
15
"""Common functionalities used by both Keras and Estimator implementations."""
16

17
18
19
20
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

21
import json
22
23
import os

24
from absl import flags
25
from absl import logging
26
import numpy as np
27
28
import tensorflow as tf

29
from official.common import distribute_utils
30
from official.recommendation import constants as rconst
31
from official.recommendation import data_pipeline
32
from official.recommendation import data_preprocessing
33
from official.recommendation import movielens
34
from official.utils.flags import core as flags_core
35

Reed's avatar
Reed committed
36
37
38
FLAGS = flags.FLAGS


Shining Sun's avatar
Shining Sun committed
39
40
41
42
def get_inputs(params):
  """Returns some parameters used by the model."""
  if FLAGS.download_if_missing and not FLAGS.use_synthetic_data:
    movielens.download(FLAGS.dataset, FLAGS.data_dir)
43

Shining Sun's avatar
Shining Sun committed
44
45
  if FLAGS.seed is not None:
    np.random.seed(FLAGS.seed)
46

Shining Sun's avatar
Shining Sun committed
47
48
  if FLAGS.use_synthetic_data:
    producer = data_pipeline.DummyConstructor()
49
    num_users, num_items = movielens.DATASET_TO_NUM_USERS_AND_ITEMS[
Shining Sun's avatar
Shining Sun committed
50
51
52
        FLAGS.dataset]
    num_train_steps = rconst.SYNTHETIC_BATCHES_PER_EPOCH
    num_eval_steps = rconst.SYNTHETIC_BATCHES_PER_EPOCH
53
  else:
Shining Sun's avatar
Shining Sun committed
54
    num_users, num_items, producer = data_preprocessing.instantiate_pipeline(
Hongkun Yu's avatar
Hongkun Yu committed
55
56
57
        dataset=FLAGS.dataset,
        data_dir=FLAGS.data_dir,
        params=params,
Shining Sun's avatar
Shining Sun committed
58
59
        constructor_type=FLAGS.constructor_type,
        deterministic=FLAGS.seed is not None)
60
61
    num_train_steps = producer.train_batches_per_epoch
    num_eval_steps = producer.eval_batches_per_epoch
62

Shining Sun's avatar
Shining Sun committed
63
  return num_users, num_items, num_train_steps, num_eval_steps, producer
64
65
66


def parse_flags(flags_obj):
Taylor Robie's avatar
Taylor Robie committed
67
  """Convenience function to turn flags into params."""
68
69
  num_gpus = flags_core.get_num_gpus(flags_obj)

70
  batch_size = flags_obj.batch_size
Taylor Robie's avatar
Taylor Robie committed
71
  eval_batch_size = flags_obj.eval_batch_size or flags_obj.batch_size
72
73
74

  return {
      "train_epochs": flags_obj.train_epochs,
75
      "batches_per_step": 1,
76
77
78
79
80
81
82
83
84
      "use_seed": flags_obj.seed is not None,
      "batch_size": batch_size,
      "eval_batch_size": eval_batch_size,
      "learning_rate": flags_obj.learning_rate,
      "mf_dim": flags_obj.num_factors,
      "model_layers": [int(layer) for layer in flags_obj.layers],
      "mf_regularization": flags_obj.mf_regularization,
      "mlp_reg_layers": [float(reg) for reg in flags_obj.mlp_regularization],
      "num_neg": flags_obj.num_neg,
85
      "distribution_strategy": flags_obj.distribution_strategy,
86
87
88
89
90
91
92
93
94
      "num_gpus": num_gpus,
      "use_tpu": flags_obj.tpu is not None,
      "tpu": flags_obj.tpu,
      "tpu_zone": flags_obj.tpu_zone,
      "tpu_gcp_project": flags_obj.tpu_gcp_project,
      "beta1": flags_obj.beta1,
      "beta2": flags_obj.beta2,
      "epsilon": flags_obj.epsilon,
      "match_mlperf": flags_obj.ml_perf,
Yuefeng Zhou's avatar
Yuefeng Zhou committed
95
      "epochs_between_evals": flags_obj.epochs_between_evals,
96
      "keras_use_ctl": flags_obj.keras_use_ctl,
97
      "hr_threshold": flags_obj.hr_threshold,
98
      "stream_files": flags_obj.tpu is not None,
99
100
101
      "train_dataset_path": flags_obj.train_dataset_path,
      "eval_dataset_path": flags_obj.eval_dataset_path,
      "input_meta_data_path": flags_obj.input_meta_data_path,
102
  }
103
104


105
def get_v1_distribution_strategy(params):
Shining Sun's avatar
Shining Sun committed
106
107
108
  """Returns the distribution strategy to use."""
  if params["use_tpu"]:
    # Some of the networking libraries are quite chatty.
Hongkun Yu's avatar
Hongkun Yu committed
109
110
111
112
    for name in [
        "googleapiclient.discovery", "googleapiclient.discovery_cache",
        "oauth2client.transport"
    ]:
Shining Sun's avatar
Shining Sun committed
113
      logging.getLogger(name).setLevel(logging.ERROR)
114

115
    tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(
Shining Sun's avatar
Shining Sun committed
116
117
118
        tpu=params["tpu"],
        zone=params["tpu_zone"],
        project=params["tpu_gcp_project"],
Hongkun Yu's avatar
Hongkun Yu committed
119
        coordinator_name="coordinator")
120

121
    logging.info("Issuing reset command to TPU to ensure a clean state.")
Shining Sun's avatar
Shining Sun committed
122
    tf.Session.reset(tpu_cluster_resolver.get_master())
123

Shining Sun's avatar
Shining Sun committed
124
125
126
127
    # Estimator looks at the master it connects to for MonitoredTrainingSession
    # by reading the `TF_CONFIG` environment variable, and the coordinator
    # is used by StreamingFilesDataset.
    tf_config_env = {
Hongkun Yu's avatar
Hongkun Yu committed
128
129
130
131
132
133
        "session_master":
            tpu_cluster_resolver.get_master(),
        "eval_session_master":
            tpu_cluster_resolver.get_master(),
        "coordinator":
            tpu_cluster_resolver.cluster_spec().as_dict()["coordinator"]
Shining Sun's avatar
Shining Sun committed
134
    }
Haoyu Zhang's avatar
Haoyu Zhang committed
135
    os.environ["TF_CONFIG"] = json.dumps(tf_config_env)
136

137
    distribution = tf.distribute.TPUStrategy(
Shining Sun's avatar
Shining Sun committed
138
        tpu_cluster_resolver, steps_per_run=100)
139

Shining Sun's avatar
Shining Sun committed
140
  else:
141
    distribution = distribute_utils.get_distribution_strategy(
Shining Sun's avatar
Shining Sun committed
142
        num_gpus=params["num_gpus"])
143

Shining Sun's avatar
Shining Sun committed
144
  return distribution
145

146
147
148
149

def define_ncf_flags():
  """Add flags for running ncf_main."""
  # Add common flags
Hongkun Yu's avatar
Hongkun Yu committed
150
151
152
153
154
155
156
157
158
159
  flags_core.define_base(
      model_dir=True,
      clean=True,
      train_epochs=True,
      epochs_between_evals=True,
      export_dir=False,
      run_eagerly=True,
      stop_threshold=True,
      num_gpu=True,
      distribution_strategy=True)
160
  flags_core.define_performance(
161
      synthetic_data=True,
Nimit Nigania's avatar
Nimit Nigania committed
162
      dtype=True,
163
      fp16_implementation=True,
Nimit Nigania's avatar
Nimit Nigania committed
164
165
      loss_scale=True,
      dynamic_loss_scale=True,
166
      enable_xla=True,
167
  )
168
  flags_core.define_device(tpu=True)
169
170
171
172
  flags_core.define_benchmark()

  flags.adopt_module_key_flags(flags_core)

173
174
  movielens.define_flags()

175
176
177
  flags_core.set_defaults(
      model_dir="/tmp/ncf/",
      data_dir="/tmp/movielens-data/",
178
      dataset=movielens.ML_1M,
179
      train_epochs=2,
180
      batch_size=99000,
Hongkun Yu's avatar
Hongkun Yu committed
181
      tpu=None)
182
183

  # Add ncf-specific flags
184
  flags.DEFINE_boolean(
Hongkun Yu's avatar
Hongkun Yu committed
185
186
187
      name="download_if_missing",
      default=True,
      help=flags_core.help_wrap(
188
189
          "Download data to data_dir if it is not already present."))

190
  flags.DEFINE_integer(
Hongkun Yu's avatar
Hongkun Yu committed
191
192
193
      name="eval_batch_size",
      default=None,
      help=flags_core.help_wrap(
194
195
196
197
198
          "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."))

199
  flags.DEFINE_integer(
Hongkun Yu's avatar
Hongkun Yu committed
200
201
      name="num_factors",
      default=8,
202
203
204
205
      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(
Hongkun Yu's avatar
Hongkun Yu committed
206
207
      name="layers",
      default=["64", "32", "16", "8"],
208
209
210
211
212
      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(
Hongkun Yu's avatar
Hongkun Yu committed
213
214
      name="mf_regularization",
      default=0.,
215
216
217
218
219
220
      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(
Hongkun Yu's avatar
Hongkun Yu committed
221
222
      name="mlp_regularization",
      default=["0.", "0.", "0.", "0."],
223
224
225
226
227
      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(
Hongkun Yu's avatar
Hongkun Yu committed
228
229
      name="num_neg",
      default=4,
230
231
232
233
      help=flags_core.help_wrap(
          "The Number of negative instances to pair with a positive instance."))

  flags.DEFINE_float(
Hongkun Yu's avatar
Hongkun Yu committed
234
235
      name="learning_rate",
      default=0.001,
236
237
      help=flags_core.help_wrap("The learning rate."))

238
  flags.DEFINE_float(
Hongkun Yu's avatar
Hongkun Yu committed
239
240
      name="beta1",
      default=0.9,
241
242
243
      help=flags_core.help_wrap("beta1 hyperparameter for the Adam optimizer."))

  flags.DEFINE_float(
Hongkun Yu's avatar
Hongkun Yu committed
244
245
      name="beta2",
      default=0.999,
246
247
248
      help=flags_core.help_wrap("beta2 hyperparameter for the Adam optimizer."))

  flags.DEFINE_float(
Hongkun Yu's avatar
Hongkun Yu committed
249
250
      name="epsilon",
      default=1e-8,
251
252
253
      help=flags_core.help_wrap("epsilon hyperparameter for the Adam "
                                "optimizer."))

254
  flags.DEFINE_float(
Hongkun Yu's avatar
Hongkun Yu committed
255
256
      name="hr_threshold",
      default=1.0,
257
258
259
260
261
262
      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."))
263

264
  flags.DEFINE_enum(
Hongkun Yu's avatar
Hongkun Yu committed
265
266
267
268
      name="constructor_type",
      default="bisection",
      enum_values=["bisection", "materialized"],
      case_sensitive=False,
269
270
271
272
273
      help=flags_core.help_wrap(
          "Strategy to use for generating false negatives. materialized has a"
          "precompute that scales badly, but a faster per-epoch construction"
          "time and can be faster on very large systems."))

274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
  flags.DEFINE_string(
      name="train_dataset_path",
      default=None,
      help=flags_core.help_wrap("Path to training data."))

  flags.DEFINE_string(
      name="eval_dataset_path",
      default=None,
      help=flags_core.help_wrap("Path to evaluation data."))

  flags.DEFINE_string(
      name="input_meta_data_path",
      default=None,
      help=flags_core.help_wrap("Path to input meta data file."))

289
  flags.DEFINE_bool(
Hongkun Yu's avatar
Hongkun Yu committed
290
291
      name="ml_perf",
      default=False,
292
293
294
295
296
297
298
299
300
301
302
303
304
      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
305
  flags.DEFINE_bool(
Hongkun Yu's avatar
Hongkun Yu committed
306
307
      name="output_ml_perf_compliance_logging",
      default=False,
Reed's avatar
Reed committed
308
309
310
311
312
313
      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 "
Hongkun Yu's avatar
Hongkun Yu committed
314
          "caches, which is required for MLPerf compliance."))
Reed's avatar
Reed committed
315

316
  flags.DEFINE_integer(
Hongkun Yu's avatar
Hongkun Yu committed
317
318
319
      name="seed",
      default=None,
      help=flags_core.help_wrap(
320
321
          "This value will be used to seed both NumPy and TensorFlow."))

Hongkun Yu's avatar
Hongkun Yu committed
322
323
324
325
  @flags.validator(
      "eval_batch_size",
      "eval_batch_size must be at least {}".format(rconst.NUM_EVAL_NEGATIVES +
                                                   1))
326
  def eval_size_check(eval_batch_size):
Taylor Robie's avatar
Taylor Robie committed
327
328
    return (eval_batch_size is None or
            int(eval_batch_size) > rconst.NUM_EVAL_NEGATIVES)
329

330
331
332
333
  flags.DEFINE_bool(
      name="early_stopping",
      default=False,
      help=flags_core.help_wrap(
Haoyu Zhang's avatar
Haoyu Zhang committed
334
          "If True, we stop the training when it reaches hr_threshold"))
335

336
337
338
339
  flags.DEFINE_bool(
      name="keras_use_ctl",
      default=False,
      help=flags_core.help_wrap(
Haoyu Zhang's avatar
Haoyu Zhang committed
340
          "If True, we use a custom training loop for keras."))
341

Haoyu Zhang's avatar
Haoyu Zhang committed
342

Shining Sun's avatar
Shining Sun committed
343
def convert_to_softmax_logits(logits):
344
  """Convert the logits returned by the base model to softmax logits.
Shining Sun's avatar
Shining Sun committed
345

346
347
348
349
350
351
  Args:
    logits: used to create softmax.

  Returns:
    Softmax with the first column of zeros is equivalent to sigmoid.
  """
Shining Sun's avatar
Shining Sun committed
352
353
  softmax_logits = tf.concat([logits * 0, logits], axis=1)
  return softmax_logits