Utility library/script to convert saved models to tflite models. (#95)

Utility library/script to convert saved models to tflite models.

Author: mtrofin

compiler_opt/rl/policy_saver.py

@@ -24,6 +24,7 @@
 from typing import Dict, Tuple
 
 OUTPUT_SIGNATURE = 'output_spec.json'
+TFLITE_MODEL_NAME = 'model.tflite'
 
 _TYPE_CONVERSION_DICT = {
     tf.float32: 'float',
@@ -158,3 +159,44 @@ def save(self, root_dir: str):
     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)

compiler_opt/rl/policy_saver_test.py

@@ -27,6 +27,55 @@
 from compiler_opt.rl import policy_saver
 
 
+# copied from the llvm regalloc generator
+def _gen_test_model(outdir: str):
+  policy_decision_label = 'index_to_evict'
+  policy_output_spec = """
+  [
+      {
+          "logging_name": "index_to_evict",
+          "tensor_spec": {
+              "name": "StatefulPartitionedCall",
+              "port": 0,
+              "type": "int64_t",
+              "shape": [
+                  1
+              ]
+          }
+      }
+  ]
+  """
+  per_register_feature_list = ['mask']
+  num_registers = 33
+
+  def get_input_signature():
+    """Returns (time_step_spec, action_spec) for LLVM register allocation."""
+    inputs = dict(
+        (key, tf.TensorSpec(dtype=tf.int64, shape=(num_registers), name=key))
+        for key in per_register_feature_list)
+    return inputs
+
+  module = tf.Module()
+  # We have to set this useless variable in order for the TF C API to correctly
+  # intake it
+  module.var = tf.Variable(0, dtype=tf.int64)
+
+  def action(*inputs):
+    result = tf.math.argmax(
+        tf.cast(inputs[0]['mask'], tf.int32), axis=-1) + module.var
+    return {policy_decision_label: result}
+
+  module.action = tf.function()(action)
+  action = {
+      'action': module.action.get_concrete_function(get_input_signature())
+  }
+  tf.saved_model.save(module, outdir, signatures=action)
+  output_spec_path = os.path.join(outdir, 'output_spec.json')
+  with tf.io.gfile.GFile(output_spec_path, 'w') as f:
+    print(f'Writing output spec to {output_spec_path}.')
+    f.write(policy_output_spec)
+
+
 class PolicySaverTest(tf.test.TestCase):
 
   def setUp(self):
@@ -80,6 +129,18 @@ def test_save_policy(self):
           }
       }], json.loads(tf.io.gfile.GFile(output_signature_fn).read()))
 
+  def test_tflite_conversion(self):
+    sm_dir = os.path.join(self.get_temp_dir(), 'saved_model')
+    tflite_dir = os.path.join(self.get_temp_dir(), 'tflite_model')
+    _gen_test_model(sm_dir)
+    policy_saver.convert_mlgo_model(sm_dir, tflite_dir)
+    self.assertTrue(
+        tf.io.gfile.exists(
+            os.path.join(tflite_dir, policy_saver.TFLITE_MODEL_NAME)))
+    self.assertTrue(
+        tf.io.gfile.exists(
+            os.path.join(tflite_dir, policy_saver.OUTPUT_SIGNATURE)))
+
 
 if __name__ == '__main__':
   tf.test.main()