[AIR] Add TorchCheckpoint.from_state_dict
#27970
Conversation
| *, | ||
| preprocessor: Optional["Preprocessor"] = None, | ||
| ) -> "TorchCheckpoint": | ||
| """Create a :class:`~ray.air.checkpoint.Checkpoint` that stores a model state |
There was a problem hiding this comment.
First line of docstring should be 1 line please!
There was a problem hiding this comment.
I'm not sure if we should do that.
If the docstring summary is on one line, then the line is longer than 88 characters. This normally means we should shorten our summary, but the only reason the summary is long is because of the :class:ray.air.checkpoint.Checkpoint Sphinx directive.
So, we either break that line limit convention or the one-line docstring summary convention. Given that the goal of the one-line summary convention is to have short summaries, I think it's fine to wrap summaries when we use Sphinx directives.
What do you think?
There was a problem hiding this comment.
Also, I don't think there's any existing precedent in our code base. PyTorch wraps their summaries. For example, https://pytorch.org/docs/stable/_modules/torch/nn/modules/linear.html#Bilinear
There was a problem hiding this comment.
The formatting guidelines we use follow PEP 257:
Multi-line docstrings consist of a summary line just like a one-line docstring, followed by a blank line, followed by a more elaborate description. The summary line may be used by automatic indexing tools; it is important that it fits on one line and is separated from the rest of the docstring by a blank line.
https://peps.python.org/pep-0257/#multi-line-docstrings
I agree though, for long sphinx directives, there doesn't seem to be an ideal solution.
There was a problem hiding this comment.
cc @richardliaw @maxpumperla for thoughts
| """ | ||
| checkpoint = cls.from_dict({PREPROCESSOR_KEY: preprocessor, MODEL_KEY: model}) | ||
| return checkpoint | ||
| """ # noqa: E501 |
There was a problem hiding this comment.
I think so. AFAIK there's no way to wrap
`Saving and Loading Models <https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict>`_.
in a way that the line is less than 88 characters.
| checkpoint = TorchCheckpoint.from_state_dict(expected_state_dict) | ||
| actual_state_dict = checkpoint.get_model(torch.nn.Linear(1, 1)).state_dict() | ||
| assert actual_state_dict == expected_state_dict | ||
|
|
There was a problem hiding this comment.
should we add a test for from_model as well?
There was a problem hiding this comment.
@justinvyu is on it. See #27056
bveeramani
added
the
tests-ok
PyTorch recommends saving state dictionaries instead of modules, but we don't support any way to do this. Signed-off-by: Balaji Veeramani [email protected]
Signed-off-by: Balaji Veeramani [email protected]
Why are these changes needed?
PyTorch recommends saving state dictionaries instead of modules, but we don't support any way to do this.
Related issue number
Closes #28158
See also #27922 and #27971.
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.