Add support for non-persistent buffers. #37191
Conversation
💊 Build failures summary and remediationsAs of commit d68705b (more details on the Dr. CI page):
🚧 1 ongoing upstream failure:These were probably caused by upstream breakages that are not fixed yet:
Extra GitHub checks: 1 failed
This comment was automatically generated by Dr. CI (expand for details).Follow this link to opt-out of these comments for your Pull Requests.Please report bugs/suggestions on the GitHub issue tracker. This comment has been revised 14 times. |
|
@ssnl Do you have time to review this PR? Thanks |
mruberry
added
module: nn
|
|
||
| This is typically used to register a buffer that should not to be | ||
| considered a model parameter. For example, BatchNorm's ``running_mean`` | ||
| is not a parameter, but is part of the persistent state. | ||
| is not a parameter, but is part of the module's state. Buffers, by | ||
| default, are persistent and will be saved alongside parameters. This |
There was a problem hiding this comment.
Could you make this clearer that the only difference here is that non persistent buffers won't show up in the state dict load/save while the persistent ones will.
There is no other difference right?
There was a problem hiding this comment.
Done. Yes, there's no other difference.
torch/nn/modules/module.py
Outdated
| @@ -85,6 +85,7 @@ def __init__(self): | |||
| self.training = True | |||
| self._parameters = OrderedDict() | |||
| self._buffers = OrderedDict() | |||
| self._persistent_buffers = set() | |||
There was a problem hiding this comment.
nit: to make this clearer that this is not a dict like the other things, you could name this self._persistent_buffers_name_set ? Or is this getting too long?
There was a problem hiding this comment.
Done. I'm not a huge fan of appending the type name to the identifier, but don't feel strongly about it.
torch/nn/modules/module.py
Outdated
| @@ -618,6 +627,7 @@ def remove_from(*dicts): | |||
| raise AttributeError( | |||
| "cannot assign parameters before Module.__init__() call") | |||
| remove_from(self.__dict__, self._buffers, self._modules) | |||
| self._persistent_buffers.discard(name) | |||
There was a problem hiding this comment.
nit: I think you can add the discard call inside the remove_from function here to make this slightly simpler.
torch/nn/modules/module.py
Outdated
| @@ -85,6 +85,7 @@ def __init__(self): | |||
| self.training = True | |||
| self._parameters = OrderedDict() | |||
| self._buffers = OrderedDict() | |||
| self._persistent_buffers = set() | |||
There was a problem hiding this comment.
Done. I'm not a huge fan of appending the type name to the identifier, but don't feel strongly about it.
|
|
||
| This is typically used to register a buffer that should not to be | ||
| considered a model parameter. For example, BatchNorm's ``running_mean`` | ||
| is not a parameter, but is part of the persistent state. | ||
| is not a parameter, but is part of the module's state. Buffers, by | ||
| default, are persistent and will be saved alongside parameters. This |
There was a problem hiding this comment.
Done. Yes, there's no other difference.
torch/nn/modules/module.py
Outdated
| @@ -618,6 +627,7 @@ def remove_from(*dicts): | |||
| raise AttributeError( | |||
| "cannot assign parameters before Module.__init__() call") | |||
| remove_from(self.__dict__, self._buffers, self._modules) | |||
| self._persistent_buffers.discard(name) | |||
There was a problem hiding this comment.
@albanD is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
(last push was just a rebase to pytorch/master) |
There was a problem hiding this comment.
@albanD is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
Ho nice, I was going to do it manually but better that way. Thanks! |
Summary: supported since pytorch 1.6: pytorch/pytorch#37191 This will not save pixel_mean/std in checkpoints. This is consistent with our current model zoo format. Differential Revision: D26576833 fbshipit-source-id: 65472efcfd71f9383bb47a192e095315037390b5
Summary: supported since pytorch 1.6: pytorch/pytorch#37191 This will not save pixel_mean/std in checkpoints. This is consistent with our current model zoo format. Reviewed By: theschnitz Differential Revision: D26576833 fbshipit-source-id: b7143d6ef8b106e873958394f80aba75fc11d2cf
Issue: #18056