neumf_model.py 16.5 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

15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
"""Defines NeuMF model for NCF framework.

Some abbreviations used in the code base:
NeuMF: Neural Matrix Factorization
NCF: Neural Collaborative Filtering
GMF: Generalized Matrix Factorization
MLP: Multi-Layer Perceptron

GMF applies a linear kernel to model the latent feature interactions, and MLP
uses a nonlinear kernel to learn the interaction function from data. NeuMF model
is a fused model of GMF and MLP to better model the complex user-item
interactions, and unifies the strengths of linearity of MF and non-linearity of
MLP for modeling the user-item latent structures.

In NeuMF model, it allows GMF and MLP to learn separate embeddings, and combine
the two models by concatenating their last hidden layer.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

36
37
import sys

38
39
from six.moves import xrange  # pylint: disable=redefined-builtin
import tensorflow as tf
40
41
from typing import Any, Dict, Text

Shawn Wang's avatar
Shawn Wang committed
42
from official.recommendation import constants as rconst
43
from official.recommendation import movielens
Shining Sun's avatar
Shining Sun committed
44
from official.recommendation import ncf_common
Shawn Wang's avatar
Shawn Wang committed
45
from official.recommendation import stat_utils
46
47


48
def sparse_to_dense_grads(grads_and_vars):
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
  """Convert sparse gradients to dense gradients.

  All sparse gradients, which are represented as instances of tf.IndexedSlices,
  are converted to dense Tensors. Dense gradients, which are represents as
  Tensors, are unchanged.

  The purpose of this conversion is that for small embeddings, which are used by
  this model, applying dense gradients with the AdamOptimizer is faster than
  applying sparse gradients.

  Args
    grads_and_vars: A list of (gradient, variable) tuples. Each gradient can
      be a Tensor or an IndexedSlices. Tensors are unchanged, and IndexedSlices
      are converted to dense Tensors.
  Returns:
    The same list of (gradient, variable) as `grads_and_vars`, except each
    IndexedSlices gradient is converted to a Tensor.
  """

  # Calling convert_to_tensor changes IndexedSlices into Tensors, and leaves
  # Tensors unchanged.
  return [(tf.convert_to_tensor(g), v) for g, v in grads_and_vars]


73
74
def neumf_model_fn(features, labels, mode, params):
  """Model Function for NeuMF estimator."""
75
76
77
  if params.get("use_seed"):
    tf.set_random_seed(stat_utils.random_int32())

78
  users = features[movielens.USER_COLUMN]
79
  items = features[movielens.ITEM_COLUMN]
80

Shining Sun's avatar
Shining Sun committed
81
82
83
  user_input = tf.keras.layers.Input(tensor=users)
  item_input = tf.keras.layers.Input(tensor=items)
  logits = construct_model(user_input, item_input, params).output
84

85
  # Softmax with the first column of zeros is equivalent to sigmoid.
Shining Sun's avatar
Shining Sun committed
86
  softmax_logits = ncf_common.convert_to_softmax_logits(logits)
87

88
  if mode == tf.estimator.ModeKeys.EVAL:
89
    duplicate_mask = tf.cast(features[rconst.DUPLICATE_MASK], tf.float32)
Shining Sun's avatar
Shining Sun committed
90
91
92
93
94
    return _get_estimator_spec_with_metrics(
        logits,
        softmax_logits,
        duplicate_mask,
        params["num_neg"],
Reed's avatar
Reed committed
95
        params["match_mlperf"],
96
        use_tpu_spec=params["use_tpu"])
97

98
99
  elif mode == tf.estimator.ModeKeys.TRAIN:
    labels = tf.cast(labels, tf.int32)
100
    valid_pt_mask = features[rconst.VALID_POINT_MASK]
101

102
    optimizer = tf.compat.v1.train.AdamOptimizer(
103
104
105
106
107
        learning_rate=params["learning_rate"],
        beta1=params["beta1"],
        beta2=params["beta2"],
        epsilon=params["epsilon"])
    if params["use_tpu"]:
108
      optimizer = tf.compat.v1.tpu.CrossShardOptimizer(optimizer)
109

110
    loss = tf.compat.v1.losses.sparse_softmax_cross_entropy(
111
        labels=labels,
112
        logits=softmax_logits,
Hongkun Yu's avatar
Hongkun Yu committed
113
        weights=tf.cast(valid_pt_mask, tf.float32))
114

115
116
    tf.identity(loss, name="cross_entropy")

117
118
    global_step = tf.compat.v1.train.get_global_step()
    tvars = tf.compat.v1.trainable_variables()
119
120
    gradients = optimizer.compute_gradients(
        loss, tvars, colocate_gradients_with_ops=True)
121
    gradients = sparse_to_dense_grads(gradients)
122
123
    minimize_op = optimizer.apply_gradients(
        gradients, global_step=global_step, name="train")
124
    update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
125
126
127
128
129
130
131
132
    train_op = tf.group(minimize_op, update_ops)

    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

  else:
    raise NotImplementedError


Shining Sun's avatar
Shining Sun committed
133
134
135
136
def _strip_first_and_last_dimension(x, batch_size):
  return tf.reshape(x[0, :], (batch_size,))


137
138
def construct_model(user_input: tf.Tensor, item_input: tf.Tensor,
                    params: Dict[Text, Any]) -> tf.keras.Model:
139
140
141
  """Initialize NeuMF model.

  Args:
Shining Sun's avatar
Shining Sun committed
142
143
    user_input: keras input layer for users
    item_input: keras input layer for items
144
    params: Dict of hyperparameters.
145

146
147
  Raises:
    ValueError: if the first model layer is not even.
148
  Returns:
149
    model:  a keras Model for computing the logits
150
151
152
153
154
155
156
157
158
159
160
161
162
163
  """
  num_users = params["num_users"]
  num_items = params["num_items"]

  model_layers = params["model_layers"]

  mf_regularization = params["mf_regularization"]
  mlp_reg_layers = params["mlp_reg_layers"]

  mf_dim = params["mf_dim"]

  if model_layers[0] % 2 != 0:
    raise ValueError("The first layer size should be multiple of 2!")

164
165
166
  # Initializer for embedding layers
  embedding_initializer = "glorot_uniform"

167
168
169
  def mf_slice_fn(x):
    x = tf.squeeze(x, [1])
    return x[:, :mf_dim]
Shining Sun's avatar
Shining Sun committed
170

171
172
173
  def mlp_slice_fn(x):
    x = tf.squeeze(x, [1])
    return x[:, mf_dim:]
Shining Sun's avatar
Shining Sun committed
174

175
176
177
  # It turns out to be significantly more effecient to store the MF and MLP
  # embedding portions in the same table, and then slice as needed.
  embedding_user = tf.keras.layers.Embedding(
178
179
      num_users,
      mf_dim + model_layers[0] // 2,
180
181
      embeddings_initializer=embedding_initializer,
      embeddings_regularizer=tf.keras.regularizers.l2(mf_regularization),
182
183
184
      input_length=1,
      name="embedding_user")(
          user_input)
185
186

  embedding_item = tf.keras.layers.Embedding(
187
188
      num_items,
      mf_dim + model_layers[0] // 2,
189
190
      embeddings_initializer=embedding_initializer,
      embeddings_regularizer=tf.keras.regularizers.l2(mf_regularization),
191
192
193
      input_length=1,
      name="embedding_item")(
          item_input)
194
195
196

  # GMF part
  mf_user_latent = tf.keras.layers.Lambda(
Hongkun Yu's avatar
Hongkun Yu committed
197
198
      mf_slice_fn, name="embedding_user_mf")(
          embedding_user)
199
  mf_item_latent = tf.keras.layers.Lambda(
Hongkun Yu's avatar
Hongkun Yu committed
200
201
      mf_slice_fn, name="embedding_item_mf")(
          embedding_item)
202
203
204

  # MLP part
  mlp_user_latent = tf.keras.layers.Lambda(
Hongkun Yu's avatar
Hongkun Yu committed
205
206
      mlp_slice_fn, name="embedding_user_mlp")(
          embedding_user)
207
  mlp_item_latent = tf.keras.layers.Lambda(
Hongkun Yu's avatar
Hongkun Yu committed
208
209
      mlp_slice_fn, name="embedding_item_mlp")(
          embedding_item)
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229

  # Element-wise multiply
  mf_vector = tf.keras.layers.multiply([mf_user_latent, mf_item_latent])

  # Concatenation of two latent features
  mlp_vector = tf.keras.layers.concatenate([mlp_user_latent, mlp_item_latent])

  num_layer = len(model_layers)  # Number of layers in the MLP
  for layer in xrange(1, num_layer):
    model_layer = tf.keras.layers.Dense(
        model_layers[layer],
        kernel_regularizer=tf.keras.regularizers.l2(mlp_reg_layers[layer]),
        activation="relu")
    mlp_vector = model_layer(mlp_vector)

  # Concatenate GMF and MLP parts
  predict_vector = tf.keras.layers.concatenate([mf_vector, mlp_vector])

  # Final prediction layer
  logits = tf.keras.layers.Dense(
Hongkun Yu's avatar
Hongkun Yu committed
230
231
232
233
234
      1,
      activation=None,
      kernel_initializer="lecun_uniform",
      name=movielens.RATING_COLUMN)(
          predict_vector)
235
236

  # Print model topology.
237
238
  model = tf.keras.models.Model([user_input, item_input], logits)
  model.summary()
239
240
  sys.stdout.flush()

241
  return model
242
243


244
245
246
247
248
249
def _get_estimator_spec_with_metrics(logits: tf.Tensor,
                                     softmax_logits: tf.Tensor,
                                     duplicate_mask: tf.Tensor,
                                     num_training_neg: int,
                                     match_mlperf: bool = False,
                                     use_tpu_spec: bool = False):
Shining Sun's avatar
Shining Sun committed
250
251
252
253
254
255
256
257
258
259
  """Returns a EstimatorSpec that includes the metrics."""
  cross_entropy, \
  metric_fn, \
  in_top_k, \
  ndcg, \
  metric_weights = compute_eval_loss_and_metrics_helper(
      logits,
      softmax_logits,
      duplicate_mask,
      num_training_neg,
260
      match_mlperf)
Shining Sun's avatar
Shining Sun committed
261
262

  if use_tpu_spec:
263
    return tf.estimator.tpu.TPUEstimatorSpec(
Shining Sun's avatar
Shining Sun committed
264
265
266
267
268
269
270
        mode=tf.estimator.ModeKeys.EVAL,
        loss=cross_entropy,
        eval_metrics=(metric_fn, [in_top_k, ndcg, metric_weights]))

  return tf.estimator.EstimatorSpec(
      mode=tf.estimator.ModeKeys.EVAL,
      loss=cross_entropy,
Hongkun Yu's avatar
Hongkun Yu committed
271
      eval_metric_ops=metric_fn(in_top_k, ndcg, metric_weights))
Shining Sun's avatar
Shining Sun committed
272
273


274
275
276
277
278
def compute_eval_loss_and_metrics_helper(logits: tf.Tensor,
                                         softmax_logits: tf.Tensor,
                                         duplicate_mask: tf.Tensor,
                                         num_training_neg: int,
                                         match_mlperf: bool = False):
279
280
281
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
  """Model evaluation with HR and NDCG metrics.

  The evaluation protocol is to rank the test interacted item (truth items)
  among the randomly chosen 999 items that are not interacted by the user.
  The performance of the ranked list is judged by Hit Ratio (HR) and Normalized
  Discounted Cumulative Gain (NDCG).

  For evaluation, the ranked list is truncated at 10 for both metrics. As such,
  the HR intuitively measures whether the test item is present on the top-10
  list, and the NDCG accounts for the position of the hit by assigning higher
  scores to hits at top ranks. Both metrics are calculated for each test user,
  and the average scores are reported.

  If `match_mlperf` is True, then the HR and NDCG computations are done in a
  slightly unusual way to match the MLPerf reference implementation.
  Specifically, if the evaluation negatives contain duplicate items, it will be
  treated as if the item only appeared once. Effectively, for duplicate items in
  a row, the predicted score for all but one of the items will be set to
  -infinity

  For example, suppose we have that following inputs:
  logits_by_user:     [[ 2,  3,  3],
                       [ 5,  4,  4]]

  items_by_user:     [[10, 20, 20],
                      [30, 40, 40]]

  # Note: items_by_user is not explicitly present. Instead the relevant \
          information is contained within `duplicate_mask`

  top_k: 2

  Then with match_mlperf=True, the HR would be 2/2 = 1.0. With
  match_mlperf=False, the HR would be 1/2 = 0.5. This is because each user has
  predicted scores for only 2 unique items: 10 and 20 for the first user, and 30
  and 40 for the second. Therefore, with match_mlperf=True, it's guaranteed the
  first item's score is in the top 2. With match_mlperf=False, this function
  would compute the first user's first item is not in the top 2, because item 20
  has a higher score, and item 20 occurs twice.

  Args:
320
321
322
    logits: A tensor containing the predicted logits for each user. The shape of
      logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits for a
      user are grouped, and the last element of the group is the true element.
323
    softmax_logits: The same tensor, but with zeros left-appended.
324
325
326
    duplicate_mask: A vector with the same shape as logits, with a value of 1 if
      the item corresponding to the logit at that position has already appeared
      for that user.
327
328
329
330
    num_training_neg: The number of negatives per positive during training.
    match_mlperf: Use the MLPerf reference convention for computing rank.

  Returns:
Shining Sun's avatar
Shining Sun committed
331
332
333
334
335
    cross_entropy: the loss
    metric_fn: the metrics function
    in_top_k: hit rate metric
    ndcg: ndcg metric
    metric_weights: metric weights
336
  """
337
338
  in_top_k, ndcg, metric_weights, logits_by_user = compute_top_k_and_ndcg(
      logits, duplicate_mask, match_mlperf)
339
340
341

  # Examples are provided by the eval Dataset in a structured format, so eval
  # labels can be reconstructed on the fly.
Hongkun Yu's avatar
Hongkun Yu committed
342
343
344
345
346
347
348
  eval_labels = tf.reshape(
      shape=(-1,),
      tensor=tf.one_hot(
          tf.zeros(shape=(logits_by_user.shape[0],), dtype=tf.int32) +
          rconst.NUM_EVAL_NEGATIVES,
          logits_by_user.shape[1],
          dtype=tf.int32))
349
350
351
352
353
354
355
356

  eval_labels_float = tf.cast(eval_labels, tf.float32)

  # During evaluation, the ratio of negatives to positives is much higher
  # than during training. (Typically 999 to 1 vs. 4 to 1) By adjusting the
  # weights for the negative examples we compute a loss which is consistent with
  # the training data. (And provides apples-to-apples comparison)
  negative_scale_factor = num_training_neg / rconst.NUM_EVAL_NEGATIVES
Hongkun Yu's avatar
Hongkun Yu committed
357
358
359
  example_weights = ((eval_labels_float +
                      (1 - eval_labels_float) * negative_scale_factor) *
                     (1 + rconst.NUM_EVAL_NEGATIVES) / (1 + num_training_neg))
360
361

  # Tile metric weights back to logit dimensions
Hongkun Yu's avatar
Hongkun Yu committed
362
363
364
  expanded_metric_weights = tf.reshape(
      tf.tile(metric_weights[:, tf.newaxis],
              (1, rconst.NUM_EVAL_NEGATIVES + 1)), (-1,))
365
366
367
368

  # ignore padded examples
  example_weights *= tf.cast(expanded_metric_weights, tf.float32)

369
  cross_entropy = tf.compat.v1.losses.sparse_softmax_cross_entropy(
370
371
372
373
      logits=softmax_logits, labels=eval_labels, weights=example_weights)

  def metric_fn(top_k_tensor, ndcg_tensor, weight_tensor):
    return {
Hongkun Yu's avatar
Hongkun Yu committed
374
375
376
377
378
379
380
381
382
        rconst.HR_KEY:
            tf.compat.v1.metrics.mean(
                top_k_tensor, weights=weight_tensor,
                name=rconst.HR_METRIC_NAME),
        rconst.NDCG_KEY:
            tf.compat.v1.metrics.mean(
                ndcg_tensor,
                weights=weight_tensor,
                name=rconst.NDCG_METRIC_NAME)
383
384
    }

Shining Sun's avatar
Shining Sun committed
385
  return cross_entropy, metric_fn, in_top_k, ndcg, metric_weights
386
387


388
389
390
def compute_top_k_and_ndcg(logits: tf.Tensor,
                           duplicate_mask: tf.Tensor,
                           match_mlperf: bool = False):
391
392
393
  """Compute inputs of metric calculation.

  Args:
394
395
396
397
398
399
    logits: A tensor containing the predicted logits for each user. The shape of
      logits is (num_users_per_batch * (1 + NUM_EVAL_NEGATIVES),) Logits for a
      user are grouped, and the first element of the group is the true element.
    duplicate_mask: A vector with the same shape as logits, with a value of 1 if
      the item corresponding to the logit at that position has already appeared
      for that user.
400
401
402
403
404
405
406
407
    match_mlperf: Use the MLPerf reference convention for computing rank.

  Returns:
    is_top_k, ndcg and weights, all of which has size (num_users_in_batch,), and
    logits_by_user which has size
    (num_users_in_batch, (rconst.NUM_EVAL_NEGATIVES + 1)).
  """
  logits_by_user = tf.reshape(logits, (-1, rconst.NUM_EVAL_NEGATIVES + 1))
408
409
  duplicate_mask_by_user = tf.cast(
      tf.reshape(duplicate_mask, (-1, rconst.NUM_EVAL_NEGATIVES + 1)),
410
      logits_by_user.dtype)
411
412
413
414
415
416
417
418
419

  if match_mlperf:
    # Set duplicate logits to the min value for that dtype. The MLPerf
    # reference dedupes during evaluation.
    logits_by_user *= (1 - duplicate_mask_by_user)
    logits_by_user += duplicate_mask_by_user * logits_by_user.dtype.min

  # Determine the location of the first element in each row after the elements
  # are sorted.
Hongkun Yu's avatar
Hongkun Yu committed
420
  sort_indices = tf.argsort(logits_by_user, axis=1, direction="DESCENDING")
421
422
423
424
425
426

  # Use matrix multiplication to extract the position of the true item from the
  # tensor of sorted indices. This approach is chosen because both GPUs and TPUs
  # perform matrix multiplications very quickly. This is similar to np.argwhere.
  # However this is a special case because the target will only appear in
  # sort_indices once.
Hongkun Yu's avatar
Hongkun Yu committed
427
428
  one_hot_position = tf.cast(
      tf.equal(sort_indices, rconst.NUM_EVAL_NEGATIVES), tf.int32)
429
  sparse_positions = tf.multiply(
Hongkun Yu's avatar
Hongkun Yu committed
430
431
      one_hot_position,
      tf.range(logits_by_user.shape[1])[tf.newaxis, :])
432
433
434
  position_vector = tf.reduce_sum(sparse_positions, axis=1)

  in_top_k = tf.cast(tf.less(position_vector, rconst.TOP_K), tf.float32)
Hongkun Yu's avatar
Hongkun Yu committed
435
  ndcg = tf.math.log(2.) / tf.math.log(tf.cast(position_vector, tf.float32) + 2)
436
437
438
  ndcg *= in_top_k

  # If a row is a padded row, all but the first element will be a duplicate.
Hongkun Yu's avatar
Hongkun Yu committed
439
440
  metric_weights = tf.not_equal(
      tf.reduce_sum(duplicate_mask_by_user, axis=1), rconst.NUM_EVAL_NEGATIVES)
441
442

  return in_top_k, ndcg, metric_weights, logits_by_user