-
Notifications
You must be signed in to change notification settings - Fork 24.4k
reduce memory usage for centered rmsprop #24170
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
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks good to me. I can merge after the comment is answered.
@@ -88,7 +88,7 @@ def step(self, closure=None): | |||
if group['centered']: | |||
grad_avg = state['grad_avg'] | |||
grad_avg.mul_(alpha).add_(1 - alpha, grad) | |||
avg = square_avg.addcmul(-1, grad_avg, grad_avg).sqrt().add_(group['eps']) | |||
avg = square_avg.addcmul(-1, grad_avg, grad_avg).sqrt_().add_(group['eps']) | |||
else: | |||
avg = square_avg.sqrt().add_(group['eps']) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We can use sqrt_
here too, is that correct?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
No. Since square_avg
is a state that we maintain, we should not modify it here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Right, good point.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@vincentqb has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
@vincentqb merged this pull request in bd054e7. |
Reduce gpu memory usage by using in-place operation