Add Export for TF backend #692

nkovela1 · 2023-08-10T16:59:55Z

This PR adapts the export library from tf.keras to Keras Core, along with adapting its corresponding test suite.
Note that this PR only implements export for the TF backend and expansion to JAX is still in the planning stage.

I have included some TODOs for further test coverage or debugging for certain edge use cases.

…heck

fchollet

Thanks for the PR!

fchollet · 2023-08-10T19:39:17Z

keras_core/backend/tensorflow/utils.py

+import tensorflow as tf
+
+
+def get_tensor_spec(t, dynamic_batch=False, name=None):


We should only have one tf utils file. We already have utils/tf_utils.py. We need to consolidate in one of them. I think backend/tf_utils.py is probably the best choice: 1. explicit name (in general avoid utils.py it's too generic), 2. confined to the TF backend folder.

Good idea, I've merged them into tf_utils.py under the TF backend folder.

fchollet · 2023-08-10T19:39:53Z

keras_core/export/export_lib.py

@@ -0,0 +1,570 @@
+"""Library for exporting inference-only Keras models/layers."""
+
+import tensorflow as tf


from keras_core.module_utils import tensorflow as tf

Fixed, thanks.

fchollet · 2023-08-10T19:40:19Z

keras_core/export/export_lib.py

+
+
+@keras_core_export("keras_core.export.ExportArchive")
+class ExportArchive(tf.__internal__.tracking.AutoTrackable):


Since it is fundamentally SavedModel specific I wonder if we should rename it to something that makes it explicit.

Alternatively we can have 1. backend-agnostic ExportArchive used for configuring endpoints, a TFSavedModelExportArchive subclass that is TF specific.

I think we can rename it once we either have a backend-agnostic abstraction or decide how we want to tie in the JAX ExportArchive. Let's discuss further on chat after this.

Is the class even reusable with something other than SavedModel in the future? It seems pretty baked into the APIs (e.g. TensorSpecs, tf.function)

We have to come up with our own flow separate from this to allow export to ONNX or other formats, since these APIs are essential for the functionality of SavedModel. It's harder to create an abstraction here that could be open-ended IMO.

There are existing TF to ONNX model converters like this one: https://github.com/onnx/tensorflow-onnx
Perhaps we can integrate some of their logic later or leave it to an OSS contribution?

fchollet · 2023-08-10T19:42:43Z

keras_core/export/export_lib_test.py

+            atol=1e-6,
+        )
+
+    # def test_model_with_lookup_table(self):


What was the issue?

The issue was calling a model with a TextVectorization Layer with string tensors (as we try to standardize them to float upon model call): https://chat.google.com/room/AAAADD9-Dbs/mZnA6hEpffQ

Since JAX and Torch do not support string tensors, we're trying a workaround of calling the layer first as Matt suggested (but there are some strange SavedModel tracked resource errors that need to be resolved with this workaround -- we can revisit this in a separate PR).

^ Disregard that

Thanks for the workaround using the Input Layer, I've added the test back, changed some of the Lookup layer export logic and it works now!

Interesting, running into another one of those errors where a test fails in CI but passes using pytest directly on it.

fchollet · 2023-08-10T19:44:00Z

keras_core/models/model.py

-        )
+        """Create a SavedModel artifact for inference (e.g. via TF-Serving).
+
+        This method lets you export a model to a lightweight SavedModel artifact


We should first mention that it only works with the TF backend.

Fixed, thanks

fchollet · 2023-08-10T22:29:38Z

keras_core/export/export_lib.py

+    export_archive.write_out(filepath)
+
+
+class ReloadedLayer(Layer):


Upon reflection, I think this is a pretty unfortunate name. I wonder if we should rename it to TFSMLayer or something to that extent? In the future we could have other formats as well, like ONNXLayer. It would make sense since it can work with more than what's saved with ExportArchive -- any SavedModel will do.

To note, this object doesn't seem to be exported to the public API yet.

Good idea, I've changed it to TFSMLayer, thanks.

fchollet · 2023-08-10T22:30:44Z

keras_core/export/export_lib.py

+
+
+@keras_core_export("keras_core.export.ExportArchive")
+class ExportArchive(tf.__internal__.tracking.AutoTrackable):


Is the class even reusable with something other than SavedModel in the future? It seems pretty baked into the APIs (e.g. TensorSpecs, tf.function)

fchollet · 2023-08-10T22:31:27Z

keras_core/models/model.py

+        """[TF backend only]* Create a TF SavedModel artifact for inference
+        (e.g. via TF-Serving).
+
+        *Note: This can currently only be used with the TF backend.


**Note:**

fchollet · 2023-08-11T01:12:14Z

The lookup table test is failing, though it isn't entirely clear what the issue is.

nkovela1 · 2023-08-11T20:34:35Z

I've decided to comment out the lookup test for now, as it's unclear how CI is affected after investigating. It's a blocker for this PR and can be addressed separately.

fchollet

LGTM, thank you!

nkovela1 and others added 28 commits July 17, 2023 15:16

Add saved model test

88e9b30

Merge remote-tracking branch 'refs/remotes/origin/main'

dfd404f

Add TF tracking attribute

19b0e39

Add tests for functional and subclassed

0be8fcc

Fix saving trackables

8908273

Fix test assertions

0418c60

Fix formatting

82c3af1

Add comments for attribute tracking

6c8731d

Merge branch 'keras-team:main' into main

ac35c30

Merge remote-tracking branch 'refs/remotes/origin/main'

d341600

Change saved model test description

8c1d954

Add backend conditional for attribute

9751d02

Change package name

c1391cb

Change epoch nums

1e7df16

Revert epochs

51410fe

Add set verbose logging utility and debug callback tests

1f11c1a

Fix formatting

e93a4a6

Sync with main repo

99301e1

Merge remote-tracking branch 'refs/remotes/origin/main'

a6eda55

Merge remote-tracking branch 'refs/remotes/origin/main'

2902b72

Merge remote-tracking branch 'refs/remotes/origin/main'

49b74b8

Initial port of model export

1f446c0

Fix imports

adbc885

Add save spec methods to TF layer

ede9bd7

Add export function to Keras Core base model

d4f2b1a

Downgrade naming error to warning and debug TF variable collections c…

766b9f6

…heck

Simplify weight reloading

b6990f8

Fix formatting, add TODOs

316fdc7

fchollet reviewed Aug 10, 2023

View reviewed changes

Unify tf_utils under backend/tensorflow

02df6af

nkovela1 added 3 commits August 10, 2023 21:13

Fix docstring and import

8f3f3c9

Fix module utils import

82bf3a3

Fix lookup layers export and add test

9fdb0dc

fchollet reviewed Aug 10, 2023

View reviewed changes

Change naming to TFSMLayer

796b466

nkovela1 added 2 commits August 11, 2023 17:50

Remove parameterized

db80dc9

Comment out failing test

4861175

fchollet approved these changes Aug 11, 2023

View reviewed changes

fchollet merged commit 138157b into keras-team:main Aug 11, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Export for TF backend #692

Add Export for TF backend #692

nkovela1 commented Aug 10, 2023

fchollet left a comment

fchollet Aug 10, 2023

nkovela1 Aug 10, 2023

fchollet Aug 10, 2023

nkovela1 Aug 10, 2023

fchollet Aug 10, 2023

nkovela1 Aug 10, 2023

fchollet Aug 10, 2023

nkovela1 Aug 11, 2023

fchollet Aug 10, 2023

nkovela1 Aug 10, 2023

nkovela1 Aug 10, 2023

nkovela1 Aug 10, 2023

fchollet Aug 10, 2023

nkovela1 Aug 10, 2023

fchollet Aug 10, 2023

nkovela1 Aug 11, 2023

fchollet Aug 10, 2023

fchollet Aug 10, 2023

nkovela1 Aug 11, 2023

fchollet commented Aug 11, 2023 •

edited

Loading

nkovela1 commented Aug 11, 2023

fchollet left a comment

		import tensorflow as tf


		def get_tensor_spec(t, dynamic_batch=False, name=None):

		@@ -0,0 +1,570 @@
		"""Library for exporting inference-only Keras models/layers."""

		import tensorflow as tf



		@keras_core_export("keras_core.export.ExportArchive")
		class ExportArchive(tf.__internal__.tracking.AutoTrackable):

		export_archive.write_out(filepath)


		class ReloadedLayer(Layer):

Add Export for TF backend #692

Add Export for TF backend #692

Conversation

nkovela1 commented Aug 10, 2023

fchollet left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fchollet commented Aug 11, 2023 • edited Loading

nkovela1 commented Aug 11, 2023

fchollet left a comment

Choose a reason for hiding this comment

fchollet commented Aug 11, 2023 •

edited

Loading