Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for VariablePolicy to saved_model.load #54319

Merged
merged 10 commits into from Jun 27, 2022
Merged

Add support for VariablePolicy to saved_model.load #54319

merged 10 commits into from Jun 27, 2022

Conversation

wenscarl
Copy link
Contributor

@wenscarl wenscarl commented Feb 10, 2022 •
edited

This PR adds python level support for VariablePolicy to saved_model.load when soft placement is disabled. A related issue is 53743 .

@google-ml-butler google-ml-butler bot added the size:S CL Change Size: Small label Feb 10, 2022
@wenscarl wenscarl mentioned this pull request Feb 10, 2022
Closed
@Saduf2019 Saduf2019 linked an issue Feb 10, 2022 that may be closed by this pull request
Python API of saved_model.load to support tf.saved_model.experimental.VariablePolicy #53743
Closed
@gbaned gbaned added this to Assigned Reviewer in PR Queue via automation Feb 10, 2022
@gbaned gbaned requested a review from dellis23 February 10, 2022 09:28
@google-ml-butler google-ml-butler bot added the awaiting review Pull request awaiting review label Feb 10, 2022
@rohan100jain rohan100jain requested review from ccrusius and removed request for rohan100jain March 2, 2022 22:58
@gbaned gbaned requested a review from kenfranko March 23, 2022 15:45
@gbaned
Copy link
Contributor

gbaned commented Apr 19, 2022

@ccrusius Can you please review this PR ? Thank you!

k-w-w
k-w-w suggested changes May 5, 2022
View reviewed changes
Copy link
Contributor

@k-w-w k-w-w left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the change! Left a few comments and I will ask the team about the API changes.

tensorflow/python/saved_model/load.py Outdated Show resolved Hide resolved
tensorflow/python/saved_model/load.py
load_with_device = (
self._save_options.experimental_variable_policy
._save_variable_devices()
and config.get_soft_device_placement()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does soft placement matter here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To make the behavior to be consistent if config.get_soft_device_placement() is disabled. An error is raised when an Op cannot be placed onto its intended device. So just to avoid sneakily fallbacking to CPU.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense

tensorflow/python/saved_model/load_options.py Show resolved Hide resolved
PR Queue automation moved this from Assigned Reviewer to Reviewer Requested Changes May 5, 2022
@k-w-w
Copy link
Contributor

k-w-w commented May 6, 2022

Additionally, could you add tests?

@wenscarl wenscarl mentioned this pull request May 11, 2022
Closed
@wenscarl
Copy link
Contributor Author

Since the tests in keras/tests/ have just been removed to keras_team/keras, a unittest is added there, check PR16525.

@wenscarl wenscarl requested a review from k-w-w May 11, 2022 21:38
@wenscarl wenscarl force-pushed the saved_model_load_variable_policy branch from c908077 to 5c90c08 Compare May 16, 2022 18:16
k-w-w
k-w-w reviewed May 17, 2022
View reviewed changes
tensorflow/python/saved_model/load.py
load_with_device = (
self._save_options.experimental_variable_policy
._save_variable_devices()
and config.get_soft_device_placement()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense

tensorflow/python/saved_model/load_test.py Outdated
@@ -47,6 +49,9 @@
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.framework import versions
from tensorflow.python.keras import Model
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tensorflow/python/keras is no longer maintained. Could you change the test so the objects inherit from the low-level model (tf.Module) instead?

@wenscarl wenscarl force-pushed the saved_model_load_variable_policy branch from 8eb251c to ec9900c Compare May 18, 2022 13:27
@wenscarl wenscarl requested a review from k-w-w May 18, 2022 14:20
k-w-w
k-w-w reviewed May 20, 2022
View reviewed changes
tensorflow/python/saved_model/load_test.py Outdated
.SAVE_VARIABLE_DEVICES
),
)
loaded = load.load(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A test that runs subprocess raises flags in my head. I would drop the OOM part of the test, and instead validate that this variable has been placed on the GPU as expected.

k-w-w
k-w-w approved these changes May 20, 2022
View reviewed changes
Copy link
Contributor

@k-w-w k-w-w left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@tensorflowbutler tensorflowbutler removed the awaiting review Pull request awaiting review label May 26, 2022
@gbaned gbaned requested review from k-w-w and removed request for k-w-w May 27, 2022 11:53
@google-ml-butler google-ml-butler bot added the awaiting review Pull request awaiting review label May 27, 2022
@k-w-w k-w-w added the kokoro:force-run Tests on submitted change label May 31, 2022
@kokoro-team kokoro-team removed the kokoro:force-run Tests on submitted change label May 31, 2022
@google-ml-butler google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Jun 13, 2022
@kokoro-team kokoro-team removed the kokoro:force-run Tests on submitted change label Jun 13, 2022
@google-ml-butler google-ml-butler bot removed the ready to pull PR ready for merge process label Jun 13, 2022
@gbaned gbaned requested a review from k-w-w June 14, 2022 05:26
@google-ml-butler google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Jun 14, 2022
@kokoro-team kokoro-team removed the kokoro:force-run Tests on submitted change label Jun 14, 2022
@tensorflowbutler tensorflowbutler removed the awaiting review Pull request awaiting review label Jun 16, 2022
k-w-w
k-w-w reviewed Jun 23, 2022
View reviewed changes
Copy link
Contributor

@k-w-w k-w-w left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@wenscarl Hi again, the tests are failing due to the build.

Can you fix the load_options build rules by adding the :save_options dependency? https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/BUILD#L645-L651

@google-ml-butler google-ml-butler bot removed the ready to pull PR ready for merge process label Jun 23, 2022
@wenscarl wenscarl requested a review from k-w-w June 23, 2022 23:22
@google-ml-butler google-ml-butler bot added the awaiting review Pull request awaiting review label Jun 23, 2022
k-w-w
k-w-w approved these changes Jun 24, 2022
View reviewed changes
Copy link
Contributor

@k-w-w k-w-w left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@google-ml-butler google-ml-butler bot added kokoro:force-run Tests on submitted change ready to pull PR ready for merge process labels Jun 24, 2022
@kokoro-team kokoro-team removed the kokoro:force-run Tests on submitted change label Jun 24, 2022
@tensorflowbutler tensorflowbutler removed the awaiting review Pull request awaiting review label Jun 26, 2022
@copybara-service copybara-service bot merged commit a79936b into tensorflow:master Jun 27, 2022
trevor-m pushed a commit to trevor-m/tensorflow that referenced this pull request Oct 20, 2022
@wenscarl @nluehr
Add support for VariablePolicy to saved_model.load
6f99bde 
c.f., upstream PR-[54319](tensorflow#54319) and nvbug-[3460007](https://nvbugswb.nvidia.com/NvBugs5/SWBug.aspx?bugid=3460007&cmtNo=).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@k-w-w k-w-w k-w-w approved these changes

@dellis23 dellis23 Awaiting requested review from dellis23

@ccrusius ccrusius Awaiting requested review from ccrusius

@kenfranko kenfranko Awaiting requested review from kenfranko

Assignees

@gbaned gbaned

Labels
ready to pull PR ready for merge process size:S CL Change Size: Small
Projects
PR Queue
  
Approved by Reviewer
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Python API of saved_model.load to support tf.saved_model.experimental.VariablePolicy
5 participants
@wenscarl @gbaned @k-w-w @kokoro-team @tensorflowbutler