TFLite: when saving a saved model, also save the tflite model (#120)

Currently just saving it next to the saved model, so we can start
iterating using tflite models with llvm. Eventually we can get more
efficient.

Author: mtrofin

compiler_opt/rl/policy_saver.py

@@ -69,6 +69,47 @@ def _get_non_identity_op(tensor):
   return tensor
 
 
+def convert_saved_model(sm_dir: str, tflite_model_path: str):
+  """Convert a saved model to tflite.
+
+  Args:
+    sm_dir: path to the saved model to convert
+
+    tflite_model_path: desired output file path. Directory structure will
+      be created by this function, as needed.
+  """
+  tf.io.gfile.makedirs(os.path.dirname(tflite_model_path))
+  converter = tf.lite.TFLiteConverter.from_saved_model(sm_dir)
+  converter.target_spec.supported_ops = [
+      tf.lite.OpsSet.TFLITE_BUILTINS,
+  ]
+  converter.allow_custom_ops = True
+  tfl_model = converter.convert()
+  with tf.io.gfile.GFile(tflite_model_path, 'wb') as f:
+    f.write(tfl_model)
+
+
+def convert_mlgo_model(mlgo_model_dir: str, tflite_model_dir: str):
+  """Convert a mlgo saved model to mlgo tflite.
+
+  Args:
+    mlgo_model_dir: path to the mlgo saved model dir. It is expected to contain
+      the saved model files (i.e. saved_model.pb, the variables dir) and the
+      output_spec.json file
+
+    tflite_model_dir: path to a directory where the tflite model will be placed.
+      The model will be named model.tflite. Alongside it will be placed a copy
+      of the output_spec.json file.
+  """
+  tf.io.gfile.makedirs(tflite_model_dir)
+  convert_saved_model(mlgo_model_dir,
+                      os.path.join(tflite_model_dir, TFLITE_MODEL_NAME))
+
+  src_json = os.path.join(mlgo_model_dir, OUTPUT_SIGNATURE)
+  dest_json = os.path.join(tflite_model_dir, OUTPUT_SIGNATURE)
+  tf.io.gfile.copy(src_json, dest_json)
+
+
 class PolicySaver(object):
   """Object that saves policy and model config file required by inference.
 
@@ -157,46 +198,11 @@ def _write_output_signature(self, saver, path):
   def save(self, root_dir: str):
     """Writes policy and model_binding.txt to root_dir/policy_name/."""
     for policy_name, (saver, _) in self._policy_saver_dict.items():
-      self._save_policy(saver, os.path.join(root_dir, policy_name))
-      self._write_output_signature(saver, os.path.join(root_dir, policy_name))
-
-
-def convert_saved_model(sm_dir: str, tflite_model_path: str):
-  """Convert a saved model to tflite.
-
-  Args:
-    sm_dir: path to the saved model to convert
-
-    tflite_model_path: desired output file path. Directory structure will
-    be created by this function, as needed.
-  """
-  tf.io.gfile.makedirs(os.path.dirname(tflite_model_path))
-  converter = tf.lite.TFLiteConverter.from_saved_model(sm_dir)
-  converter.target_spec.supported_ops = [
-      tf.lite.OpsSet.TFLITE_BUILTINS,
-  ]
-  tfl_model = converter.convert()
-  with tf.io.gfile.GFile(tflite_model_path, 'wb') as f:
-    f.write(tfl_model)
-
-
-def convert_mlgo_model(mlgo_model_dir: str, tflite_model_dir: str):
-  """Convert a mlgo saved model to mlgo tflite.
-
-  Args:
-    mlgo_model_dir: path to the mlgo saved model dir. It is expected to contain
-    the saved model files (i.e. saved_model.pb, the variables dir) and the
-    output_spec.json file
-
-    tflite_model_dir: path to a directory where the tflite model will be placed.
-    The model will be named model.tflite. Alongside it will be placed a copy of
-    the output_spec.json file.
-  """
-  tf.io.gfile.makedirs(tflite_model_dir)
-  convert_saved_model(mlgo_model_dir,
-                      os.path.join(tflite_model_dir, TFLITE_MODEL_NAME))
-
-  json_file = 'output_spec.json'
-  src_json = os.path.join(mlgo_model_dir, json_file)
-  dest_json = os.path.join(tflite_model_dir, json_file)
-  tf.io.gfile.copy(src_json, dest_json)
+      saved_model_dir = os.path.join(root_dir, policy_name)
+      self._save_policy(saver, saved_model_dir)
+      self._write_output_signature(saver, saved_model_dir)
+      # This is not quite the most efficient way to do this - we save the model
+      # just to load it again and save it as tflite - but it's the minimum,
+      # temporary step so we can validate more thoroughly our use of tflite.
+      convert_saved_model(saved_model_dir,
+                          os.path.join(saved_model_dir, TFLITE_MODEL_NAME))

compiler_opt/rl/policy_saver_test.py

@@ -113,6 +113,8 @@ def test_save_policy(self):
     for sub_dir in ['saved_policy', 'saved_collect_policy']:
       self.assertTrue(
           tf.io.gfile.exists(os.path.join(root_dir, sub_dir, 'saved_model.pb')))
+      self.assertTrue(
+          tf.io.gfile.exists(os.path.join(root_dir, sub_dir, 'model.tflite')))
       self.assertTrue(
           tf.io.gfile.exists(
               os.path.join(root_dir, sub_dir,