Merge pull request #133 from victorromeo/api_unit_tests

Unit tests for ONNX export/convert

setup.py

@@ -47,6 +47,7 @@ class _Develop(_CompileFlatbuffMixin, _develop):
         "Jinja2",
         "tensorflow>=2.1.0",
         "onnx",
+        "keras2onnx",
         "idx2numpy",
         "attrs",
         "click",

tests/test_api/test_keras_onnx/conftest.py

@@ -0,0 +1,62 @@
+from pytest import fixture
+
+import os
+import numpy as np
+
+import tensorflow as tf
+from tensorflow.keras.models import Sequential, save_model, load_model
+from tensorflow.keras.layers import MaxPool2D, ReLU, Conv2D, Softmax, Dense, Flatten
+from tensorflow.keras.losses import sparse_categorical_crossentropy
+from tensorflow.keras.optimizers import Adam
+
+@fixture(scope='session', name='keras_model')
+def keras_model():
+
+    input_shape = (28,28,1)
+    no_classes = 10
+
+    model = Sequential()
+    model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
+    model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
+    model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
+    model.add(Flatten())
+    model.add(Dense(128, activation='relu'))
+    model.add(Dense(no_classes, activation='softmax'))
+
+    model.compile(
+        loss=sparse_categorical_crossentropy,
+        optimizer=Adam(),
+        metrics=['accuracy']
+    )
+
+    np.random.seed(12345)
+    mu, sigma = 0, 0.1 # mean and standard deviation
+    x = np.random.normal(mu, sigma, size = (1,) + input_shape)
+    y = model(x)
+
+    return model
+
+@fixture(scope='session', name='keras_model_dset')
+def keras_model_dset():
+    mnist = tf.keras.datasets.mnist
+
+    (x_train, y_train), (x_test, y_test) = mnist.load_data()
+    x_train, x_test = x_train / 255.0, x_test / 255.0
+
+    # Add a channels dimension
+    x_train = x_train[..., tf.newaxis]
+    x_test = x_test[..., tf.newaxis]
+
+    num_calibration_steps = 128
+    calibration_dtype = tf.float32
+    input_shape = (28,28,1)
+
+    def representative_dataset_gen():
+        for _ in range(num_calibration_steps):
+            rand_idx = np.random.randint(0, x_test.shape[0]-1)
+            sample = x_test[rand_idx]
+            sample = sample[tf.newaxis, ...]
+            sample = tf.cast(sample, dtype=calibration_dtype)
+            yield [sample]
+
+    return representative_dataset_gen

tests/test_api/test_keras_onnx/test_export.py

@@ -0,0 +1,35 @@
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+import utensor_cgen.api.export as export
+import tensorflow.keras as keras
+
+def test_keras_model(keras_model, keras_model_dset):
+
+    assert keras_model, 'Keras Model generation failed'
+
+    export.keras_onnx_export(
+        keras_model, 
+        representive_dataset=keras_model_dset,
+        model_name='model',
+        target='utensor'
+    )
+
+
+def test_keras_model_path(keras_model, keras_model_dset):
+  with tempfile.TemporaryDirectory(prefix='utensor_') as tmp_dir:
+    dir_path = Path(tmp_dir)
+    keras_model_path = os.path.join(dir_path, 'model.h5')
+    keras.models.save_model(
+        model=keras_model, 
+        filepath=keras_model_path
+    )
+
+    export.keras_onnx_export(
+        keras_model, 
+        representive_dataset=keras_model_dset,
+        model_name='model',
+        target='utensor'
+    )

tests/test_api/test_keras_tflite/conftest.py

@@ -0,0 +1,63 @@
+from pytest import fixture
+
+import os
+import numpy as np
+
+import tensorflow as tf
+from tensorflow.keras.models import Sequential, save_model, load_model
+from tensorflow.keras.layers import MaxPool2D, ReLU, Conv2D, Softmax, Dense, Flatten
+from tensorflow.keras.losses import sparse_categorical_crossentropy
+from tensorflow.keras.optimizers import Adam
+
+@fixture(scope='session', name='keras_model')
+def keras_model():
+
+    input_shape = (28,28,1)
+    no_classes = 10
+
+    model = Sequential()
+    model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape))
+    model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
+    model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
+    model.add(Flatten())
+    model.add(Dense(128, activation='relu'))
+    model.add(Dense(no_classes, activation='softmax'))    
+
+
+    model.compile(
+        loss=sparse_categorical_crossentropy,
+        optimizer=Adam(),
+        metrics=['accuracy']
+    )
+
+    np.random.seed(12345)
+    mu, sigma = 0, 0.1 # mean and standard deviation
+    x = np.random.normal(mu, sigma, size = (1,) + input_shape)
+    y = model(x)
+
+    return model
+
+@fixture(scope='session', name='keras_model_dset')
+def keras_model_dset():
+    mnist = tf.keras.datasets.mnist
+
+    (x_train, y_train), (x_test, y_test) = mnist.load_data()
+    x_train, x_test = x_train / 255.0, x_test / 255.0
+
+    # Add a channels dimension
+    x_train = x_train[..., tf.newaxis]
+    x_test = x_test[..., tf.newaxis]
+
+    num_calibration_steps = 128
+    calibration_dtype = tf.float32
+    input_shape = (28,28,1)
+
+    def representative_dataset_gen():
+        for _ in range(num_calibration_steps):
+            rand_idx = np.random.randint(0, x_test.shape[0]-1)
+            sample = x_test[rand_idx]
+            sample = sample[tf.newaxis, ...]
+            sample = tf.cast(sample, dtype=calibration_dtype)
+            yield [sample]
+
+    return representative_dataset_gen

tests/test_api/test_keras_tflite/test_export.py

@@ -0,0 +1,34 @@
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+import utensor_cgen.api.export as export
+import tensorflow.keras as keras
+
+def test_keras_model(keras_model, keras_model_dset):
+  assert keras_model, 'Keras Model generation failed'
+
+  export.tflm_keras_export(
+    keras_model, 
+    representive_dataset=keras_model_dset,
+    model_name='model',
+    target='utensor'
+  )
+
+def test_keras_model_path(keras_model, keras_model_dset):
+  with tempfile.TemporaryDirectory(prefix='utensor_') as tmp_dir:
+    dir_path = Path(tmp_dir)
+    keras_model_path = os.path.join(dir_path, 'model')
+    keras.models.save_model(
+        model=keras_model, 
+        filepath=keras_model_path,
+        save_format='tf'
+    )
+
+    export.tflm_keras_export(
+      keras_model_path, 
+      representive_dataset=keras_model_dset,
+      model_name='model',
+      target='utensor'
+    )

utensor_cgen/api/export.py

@@ -1,8 +1,15 @@
 import tempfile
 from pathlib import Path
 
+import os
+
 import tensorflow as tf
 
+import torch
+
+import onnx
+import keras2onnx 
+
 from .convert import convert_graph
 
 
@@ -42,3 +49,81 @@ def tflm_keras_export(
   if output_tflite_fname:
     with open(output_tflite_fname, 'wb') as fid:
       fid.write(tflm_model_content)
+
+def pytorch_onnx_export(
+  model_or_path,
+  representive_dataset,
+  model_name=None,
+  optimizations=None,
+  config_file='utensor_cli.toml',
+  target='utensor',
+  output_onnx_fname=None,
+  verbose=False
+) :
+  with tempfile.TemporaryDirectory(prefix='utensor_') as tmp_dir:
+    dir_path = Path(tmp_dir)
+    if isinstance(model_or_path, str):
+      model = torch.load(model_or_path)
+    elif isinstance(model_or_path, torch.nn.Module):
+      model = model_or_path
+    else:
+      raise RuntimeError(
+        "expecting a pytorch model or a path to torch model, get {}".format(
+          model_or_path
+        )
+      )
+
+    onnx_path = os.path.join(dir_path,f"{model_name}.onnx")
+
+    torch.onnx.export(model, 
+      representive_dataset, 
+      onnx_path, 
+      verbose=verbose
+    )
+
+    # Print a human readable representation of the graph
+    onnx_model = onnx.load(onnx_path)
+    onnx.checker.check_model(onnx_model)
+
+    if verbose:
+      onnx.helper.printable_graph(onnx_model.graph)
+
+    convert_graph(onnx_path,
+      config=config_file,
+      model_name=model_name,
+      target=target,
+    )
+
+def keras_onnx_export(
+  model_or_path,
+  representive_dataset,
+  model_name=None,
+  optimizations=None,
+  config_file='utensor_cli.toml',
+  target='utensor',
+  output_onnx_fname=None):
+
+  with tempfile.TemporaryDirectory(prefix='utensor_') as tmp_dir:
+    dir_path = Path(tmp_dir)
+
+    if isinstance(model_or_path, str):
+      model = tf.keras.model.load_model(model_or_path)
+    elif isinstance(model_or_path, tf.keras.Model):
+      model = model_or_path
+    else:
+      raise RuntimeError(
+        "expecting a keras model or a path to saved model, get {}".format(
+          model_or_path
+        )
+      )
+
+    # Perform Keras to ONNX conversion
+    onnx_model_name = f'{model_name}.onnx'
+    onnx_model = keras2onnx.convert_keras(model, model.name)
+    onnx.save_model(onnx_model, onnx_model_name)
+
+    convert_graph(onnx_model_name, 
+      config=config_file, 
+      model_name=model_name, 
+      target=target
+    )

utensor_cgen/frontend/onnx.py

@@ -27,6 +27,12 @@ def _convert_op_attribute(attrib_pb):
     return attrib_pb.i
   elif attrib_pb.HasField('s'):
     return attrib_pb.s
+  elif attrib_pb.floats:
+    return TensorProtoConverter.__utensor_generic_type__(np_array=np.array(attrib_pb.floats))
+  elif attrib_pb.ints:
+    return TensorProtoConverter.__utensor_generic_type__(np_array=np.array(attrib_pb.ints))
+  elif attrib_pb.strings:
+    return TensorProtoConverter.__utensor_generic_type__(np_array=np.array(attrib_pb.strings))
   else:
     raise ValueError('Unknown attribute value: {}'.format(attrib_pb))