Draft: Java API to use tf.function available on SavedModel.

[Shajan]

Python models that contain tf.function is inconvenient to be consumed by Java clients.
This proposal provides an API to
(a) Invoke a tf.function, given the signature name
(b) Retrieve the node name in the graph corresponding to a tf.function

TBD: Unit tests before this change is ready to merge.

[googlebot]

Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[karllessard]

Hi @Shajan , can you please rebase this PR to merge it in the shared branch created for this story we will working on together instead of master? https://github.com/tensorflow/java/tree/shared-saved-model

I think we should merge it back to master only when the whole solution is complete

[Shajan]

@karllessard i tried to rebase and then push the change to tensorflow:shared-saved-model i may not have permission to do so. Instead i made a pull request to the shared-saved-model branch.

[karllessard]

Thanks @Shajan , it's ok I've merged it and sent you an invite to collaborate to this repository, I think after you accept it, I will be able to give you write access to the shared branch

[Shajan]

Thanks @karllessard got the invite. Will wait for you to merge your changes and then add tests.

[yzhuang]

Read through this change, and I believe the following API design would be both cleaner and more intuitive to use. Wdyt?

1. Update Loader and add withSignature(String... signatures).
2. Update Loader.load() to also construct TfFunctions corresponding to the specified signatures.
3. SignatureToNodeName can become private and not exposed to end user, preferrably living inside TfFunction as private.
4. TfFunction should not have a reference to the session. Let SavedModelBundle manage the session (which it already does). TfFunction can become a pure data class.
5. Add SavedModelBundle.call(String signature, Map<String, Tensor<?>> inputs, Map<String, Tensor<?>> outputs) to SavedModelBundle. TfFunction can become private and not exposed to end user as well.

What do you think about the above proposed API design?

[Shajan]

Not exposing SignatureToNodeName (3) sounds good
TtFunction not having reference to session (4) is good as well

Regarding (5), it is desirable to have a way to do repeated call to the same function. Having a TfFuction class allows for this. I am also thinking removing runtime check currently done with each call in the Tensor call(Tensor) method and do that once.

[karllessard]

Thanks for your proposal @yzhuang , some thoughts on it:

SignatureToNodeName can become private and not exposed to end user, preferrably living inside TfFunction as private.

Looking at it, getSignatureToNodeName() can already be private in SavedModelBundle and SignatureToNodeName could be restricted at the default-package level. I don't think it should be in TfFunction though as it maps all signatures while an instance of TfFunction is only mapped to one of them.

TfFunction should not have a reference to the session. Let SavedModelBundle manage the session (which it already does). TfFunction can become a pure data class.

That would prevent though to just invoke function.call(tensor) to run the graph, which I personally like. While I understand that data classes have their lot of advantages, I'm not sure there is real gains for having TfFunction as one of them, especially that Session instances are thread-safe.

Add SavedModelBundle.call(String signature, Map<String, Tensor> inputs, Map> outputs) to SavedModelBundle. TfFunction can become private and not exposed to end user as well.

Again, I personally like the OO approach of letting users manipulating callable entities instead of having SavedModelBundle acting as a "service". Is the intention here is just to hide the TfFunction at the user level? Note that I was planning to reuse the same object to add new signatures to a model when exporting it (but I could also do it just with SavedModelBundle if we decide to take that direction).

[yzhuang]

Hi Karl and Shajan,

I have been busy with traveling, and didn't have time to follow up on this thread. Please feel free to take whatever you find helpful from my suggestions, and discard the rest. Thank you for taking time to go through my comments!

[yzhuang]

Hi @karllessard and @Shajan ,

Sorry for the late reply on your comment—just landed in Hong Kong :)

My suggestion to let SavedModel manage the session is not about thread safety, but about resource ownership. SavedModel currently “owns” the session and is responsible for closing the resource to avoid memory leaks. If we create TFFunction objects holding references to the session, it complicate the resource ownership mental model. For example, we will need to do things such as the below:

1. Perform reference counting on the session, and have the last “user” of the session close it. This is user friendly, but introduces complexity and room for bugs. OR
2. Continue to let SavedModel own the session, but TFFunction objects need to anticipate that it’s underlying resource (the session) can become defunct at anytime. This can be counterintuitive to users as users likely need to surround their call with a try/catch.

Is the intention here is just to hide the TfFunction at the user level?

My intention was to not have to think about the resource ownership problem, and keep the resource ownership unchanged. The API of having TFFunctions be callable and exposed to users sounds great to me, and we need to think about the above 1 & 2 though. This is not a contrived scenario: we already do model hot swapping at Twitter in our prediction servers, and SavedModels are hot swapped without restarting the JVM.

Thank you Karl and Shajan!

[karllessard]

@yzhuang , this SIG is having a community video chat where we will discuss about this topic, among other things. Are you available and interested to join?

[Shajan]

@karllessard missed the meeting today. I see the decision to be

Confirmation: lets model the API with functions while still loading/saving session-centric graphs

I take it that we go with the current branch (after addressing open comments) and adding unit tests.

For unit tests: Here is a proposal

• Checkin savedmodel (created using python) as a resource
• [optionally] Add the python model + instructions (but not have the test run python code)