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.

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

Clear .unconstrained weakrefs before pickling; rebuild them more often #3212

Merged
merged 3 commits into from
May 17, 2023

Conversation

fritzo
Copy link
Member

@fritzo fritzo commented May 17, 2023

Resolves #3201
Replaces #3206

This builds on @ilia-kats's diagnosis of why torch.save is failing with various Pyro objects in torch>=2, namely that the .unconstrained attributes of tensors are breaking pickling. This PR:

  • removes .unconstrained in one other place,
  • rebuilds .unconstrained in one place,
  • fixes __getstate__ in a number of places to avoid serializing weakref.refs
  • fixes one module_local_param comment.

Tested

  • ran serialization tests locally under torch==2.0.1

@fritzo fritzo requested a review from eb8680 May 17, 2023 19:53
@@ -75,3 +76,22 @@ def iter_plates_to_shape(shape):
# Go backwards (right to left)
for i, s in enumerate(shape[::-1]):
yield pyro.plate("plate_" + str(i), s)


def check_no_weakref(obj, path="", avoid_ids=None):
Copy link
Member Author

Choose a reason for hiding this comment

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

This is just a helper to debug similar issues in the future.

Copy link
Member

@eb8680 eb8680 left a comment

Choose a reason for hiding this comment

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

LGTM, thanks for tracking this down. This might also be sufficient to fix #3192

@eb8680 eb8680 merged commit 9dcd355 into dev May 17, 2023
@eb8680 eb8680 deleted the fix-unconstrained branch May 17, 2023 21:44
@eb8680 eb8680 mentioned this pull request May 17, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[bug] param store's .save() method fails with pytorch 2.0.0
2 participants