Add recommendation model spec to model maker.

PiperOrigin-RevId: 347345279

Author: lintian06

lite/examples/recommendation/ml/model/recommendation_model_launcher_keras.py

@@ -30,37 +30,40 @@
 
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string('training_data_filepattern', None,
-                    'File pattern of the training data.')
-flags.DEFINE_string('testing_data_filepattern', None,
-                    'File pattern of the training data.')
-flags.DEFINE_string('model_dir', None, 'Directory to store checkpoints.')
-flags.DEFINE_string(
-    'params_path', None,
-    'Path to the json file containing params needed to run '
-    'p13n recommendation model.')
-flags.DEFINE_integer('batch_size', 1, 'Training batch size.')
-flags.DEFINE_float('learning_rate', 0.1, 'Learning rate.')
-flags.DEFINE_integer('steps_per_epoch', 10,
-                     'Number of steps to run in each epoch.')
-flags.DEFINE_integer('num_epochs', 10000, 'Number of training epochs.')
-flags.DEFINE_integer('num_eval_steps', 1000, 'Number of eval steps.')
-flags.DEFINE_enum('run_mode', 'train_and_eval', ['train_and_eval', 'export'],
-                  'Mode of the launcher, default value is: train_and_eval')
-flags.DEFINE_float('gradient_clip_norm', 1.0,
-                   'gradient_clip_norm <= 0 meaning no clip.')
-flags.DEFINE_integer('max_history_length', 10, 'Max length of user history.')
-flags.DEFINE_integer('num_predictions', 100,
-                     'Num of top predictions to output.')
-flags.DEFINE_string(
-    'encoder_type', 'bow', 'Type of the encoder for context'
-    'encoding, the value could be ["bow", "rnn", "cnn"].')
-flags.DEFINE_string('checkpoint_path', '', 'Path to the checkpoint.')
-
 CONTEXT = 'context'
 LABEL = 'label'
 
 
+def define_flags():
+  """Define flags."""
+  flags.DEFINE_string('training_data_filepattern', None,
+                      'File pattern of the training data.')
+  flags.DEFINE_string('testing_data_filepattern', None,
+                      'File pattern of the training data.')
+  flags.DEFINE_string('model_dir', None, 'Directory to store checkpoints.')
+  flags.DEFINE_string(
+      'params_path', None,
+      'Path to the json file containing params needed to run '
+      'p13n recommendation model.')
+  flags.DEFINE_integer('batch_size', 1, 'Training batch size.')
+  flags.DEFINE_float('learning_rate', 0.1, 'Learning rate.')
+  flags.DEFINE_integer('steps_per_epoch', 10,
+                       'Number of steps to run in each epoch.')
+  flags.DEFINE_integer('num_epochs', 10000, 'Number of training epochs.')
+  flags.DEFINE_integer('num_eval_steps', 1000, 'Number of eval steps.')
+  flags.DEFINE_enum('run_mode', 'train_and_eval', ['train_and_eval', 'export'],
+                    'Mode of the launcher, default value is: train_and_eval')
+  flags.DEFINE_float('gradient_clip_norm', 1.0,
+                     'gradient_clip_norm <= 0 meaning no clip.')
+  flags.DEFINE_integer('max_history_length', 10, 'Max length of user history.')
+  flags.DEFINE_integer('num_predictions', 100,
+                       'Num of top predictions to output.')
+  flags.DEFINE_string(
+      'encoder_type', 'bow', 'Type of the encoder for context'
+      'encoding, the value could be ["bow", "rnn", "cnn"].')
+  flags.DEFINE_string('checkpoint_path', '', 'Path to the checkpoint.')
+
+
 class SimpleCheckpoint(tf.keras.callbacks.Callback):
   """Keras callback to save tf.train.Checkpoints."""
 
@@ -250,4 +253,5 @@ def main(_):
 
 
 if __name__ == '__main__':
+  define_flags()
   app.run(main)

tensorflow_examples/lite/model_maker/core/task/model_spec/__init__.py

@@ -15,7 +15,9 @@
 
 import inspect
 
+
 from tensorflow_examples.lite.model_maker.core.task.model_spec import audio_spec
+from tensorflow_examples.lite.model_maker.core.task.model_spec import recommendation_spec
 from tensorflow_examples.lite.model_maker.core.task.model_spec.image_spec import efficientnet_lite0_spec
 from tensorflow_examples.lite.model_maker.core.task.model_spec.image_spec import efficientnet_lite1_spec
 from tensorflow_examples.lite.model_maker.core.task.model_spec.image_spec import efficientnet_lite2_spec
@@ -61,6 +63,11 @@
 
     # Audio classification
     'audio_browser_fft': audio_spec.BrowserFFTSpec,
+
+    # Recommendation
+    'recommendation_bow': recommendation_spec.recommendation_bow_spec,
+    'recommendation_cnn': recommendation_spec.recommendation_cnn_spec,
+    'recommendation_rnn': recommendation_spec.recommendation_rnn_spec,
 }
 
 # List constants for supported models.
@@ -73,6 +80,11 @@
 ]
 QUESTION_ANSWERING_MODELS = ['bert_qa', 'mobilebert_qa', 'mobilebert_qa_squad']
 AUDIO_CLASSIFICATION_MODELS = ['audio_browser_fft']
+RECOMMENDATION_MODELS = [
+    'recommendation_bow',
+    'recommendation_rnn',
+    'recommendation_cnn',
+]
 
 
 def get(spec_or_str):

tensorflow_examples/lite/model_maker/core/task/model_spec/model_spec_test.py

@@ -18,12 +18,18 @@
 
 import os
 
+from absl.testing import parameterized
 import tensorflow.compat.v2 as tf
 
 from tensorflow_examples.lite.model_maker.core.task import model_spec as ms
 
+MODELS = (
+    ms.IMAGE_CLASSIFICATION_MODELS + ms.TEXT_CLASSIFICATION_MODELS +
+    ms.QUESTION_ANSWERING_MODELS + ms.AUDIO_CLASSIFICATION_MODELS +
+    ms.RECOMMENDATION_MODELS)
 
-class ModelSpecTest(tf.test.TestCase):
+
+class ModelSpecTest(tf.test.TestCase, parameterized.TestCase):
 
   def test_get(self):
     spec = ms.get('mobilenet_v2')
@@ -35,6 +41,12 @@ def test_get(self):
     spec = ms.get(ms.mobilenet_v2_spec)
     self.assertIsInstance(spec, ms.ImageModelSpec)
 
+  @parameterized.parameters(MODELS)
+  def test_get_not_none(self, model):
+    spec = ms.get(model)
+    self.assertIsNotNone(spec)
+
+  def test_get_raises(self):
     with self.assertRaises(KeyError):
       ms.get('not_exist_model_spec')
 

tensorflow_examples/lite/model_maker/core/task/model_spec/recommendation_spec.py

@@ -0,0 +1,97 @@
+# Copyright 2020 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.
+"""Recommendation model specification."""
+
+import functools
+
+import tensorflow as tf  # pylint: disable=unused-import
+
+HAS_RECOMMENDATION = True
+try:
+  from tensorflow_examples.lite.examples.recommendation.ml.model import recommendation_model as rm  # pylint: disable=g-import-not-at-top
+except ImportError:
+  HAS_RECOMMENDATION = False
+
+
+class RecommendationSpec(object):
+  """Recommendation model spec."""
+
+  def __init__(self,
+               encoder_type='bow',
+               context_embedding_dim=128,
+               label_embedding_dim=32,
+               item_vocab_size=16,
+               num_predictions=10,
+               hidden_layer_dim_ratios=None,
+               conv_num_filter_ratios=None,
+               conv_kernel_size=None,
+               lstm_num_units=None):
+    """Initialize spec.
+
+    Args:
+      encoder_type: str, encoder type. One of ('bow', 'cnn', 'rnn').
+      context_embedding_dim: int, dimension of context embedding layer.
+      label_embedding_dim: int, dimension of label embedding layer.
+      item_vocab_size: int, the size of items to be predict.
+      num_predictions: int, the number of top-K predictions in the output.
+      hidden_layer_dim_ratios: list of float, number of units in hidden layers
+        specified by ratios. default: [1.0, 0.5, 0.25].
+      conv_num_filter_ratios: list of int, for 'cnn', Conv1D layers' filter
+        ratios based on context_embedding_dim.
+      conv_kernel_size: int, for 'rnn', Conv1D layers' kernel size.
+      lstm_num_units: int, for 'rnn', LSTM layer's unit number.
+    """
+    hidden_layer_dim_ratios = hidden_layer_dim_ratios or [1.0, 0.5, 0.25]
+
+    if encoder_type == 'cnn':
+      conv_num_filter_ratios = conv_num_filter_ratios or [2, 4]
+      conv_kernel_size = conv_kernel_size or 4
+    elif encoder_type == 'rnn':
+      lstm_num_units = lstm_num_units or 16
+
+    self.encoder_type = encoder_type
+    self.context_embedding_dim = context_embedding_dim
+    self.label_embedding_dim = label_embedding_dim
+    self.hidden_layer_dim_ratios = hidden_layer_dim_ratios
+    self.item_vocab_size = item_vocab_size
+    self.num_predictions = num_predictions
+    self.conv_num_filter_ratios = conv_num_filter_ratios
+    self.conv_kernel_size = conv_kernel_size
+    self.lstm_num_units = lstm_num_units
+
+    self._params = {
+        'encoder_type': encoder_type,
+        'context_embedding_dim': context_embedding_dim,
+        'label_embedding_dim': label_embedding_dim,
+        'hidden_layer_dim_ratios': hidden_layer_dim_ratios,
+        'item_vocab_size': item_vocab_size,
+        'num_predictions': num_predictions,
+        'conv_num_filter_ratios': conv_num_filter_ratios,
+        'conv_kernel_size': conv_kernel_size,
+        'lstm_num_units': lstm_num_units,
+    }
+
+  def create_model(self):
+    """Creates recommendation model based on params."""
+    if not HAS_RECOMMENDATION:
+      return None
+    return rm.RecommendationModel(self._params)
+
+
+recommendation_bow_spec = functools.partial(
+    RecommendationSpec, encoder_type='bow')
+recommendation_cnn_spec = functools.partial(
+    RecommendationSpec, encoder_type='cnn')
+recommendation_rnn_spec = functools.partial(
+    RecommendationSpec, encoder_type='rnn')

tensorflow_examples/lite/model_maker/core/task/model_spec/recommendation_spec_test.py

@@ -0,0 +1,39 @@
+# Copyright 2020 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.
+"""Tests for recommendation spec."""
+
+from absl.testing import parameterized
+import tensorflow.compat.v2 as tf
+
+from tensorflow_examples.lite.model_maker.core.task.model_spec import recommendation_spec
+
+
+class RecommendationSpecTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('bow'),
+      ('cnn'),
+      ('rnn'),
+  )
+  def test_create_recommendation_model(self, encoder_type):
+    spec = recommendation_spec.RecommendationSpec(encoder_type)
+    model = spec.create_model()
+    if recommendation_spec.HAS_RECOMMENDATION:
+      self.assertIsInstance(model, recommendation_spec.rm.RecommendationModel)
+    else:
+      self.assertIsNone(model)
+
+
+if __name__ == '__main__':
+  tf.test.main()