Files
AMD-SHARK-Studio/tank/tflite/image_stylization_test.py
Prashant Kumar 0dcf387089 Add shark_importer for torch_models. (#183)
All the torch_models are imported to gs::shark_tank.
Scripts have been updated.
2022-07-12 20:38:19 -07:00

34 lines
1.0 KiB
Python

# RUN: %PYTHON %s
import absl.testing
import numpy
import test_util
model_path = "https://tfhub.dev/sayakpaul/lite-model/arbitrary-image-stylization-inceptionv3-dynamic-shapes/int8/predict/1?lite-format=tflite"
# Failure is due to avg_pool2d.
class ImageStylizationTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(ImageStylizationTest, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(ImageStylizationTest, self).compare_results(
iree_results, tflite_results, details
)
iree = iree_results[0].flatten().astype(numpy.single)
tflite = tflite_results[0].flatten().astype(numpy.single)
# Error is not tiny but appears close.
self.assertTrue(
numpy.isclose(
numpy.max(numpy.abs(iree - tflite)), 0.0, atol=5e-2
).all()
)
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == "__main__":
absl.testing.absltest.main()