[optimizer] refactor Adam to use functional API #44791
Conversation
[ghstack-poisoned]
💊 CI failures summary and remediationsAs of commit c3a26d6 (more details on the Dr. CI page): 💚 💚 Looks good so far! There are no failures yet. 💚 💚 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 or post in the (internal) Dr. CI Users group. This comment has been revised 51 times. |
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
wanchaol
requested review from
albanD,
izdeby,
mrshenli,
pritamdamania87 and
vincentqb
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
|
@egrefen @seba-1511 -- since you were part of the #39279 discussion, do you have any feedback on the current functional implementation? |
[ghstack-poisoned]
albanD
left a comment
albanD
left a comment
There was a problem hiding this comment.
Looks good.
Just the init that should be moved back to be lazy.
torch/optim/adam.py
Outdated
| state = self.state[p] | ||
| state['step'] = 0 | ||
| # Exponential moving average of gradient values | ||
| state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format) |
There was a problem hiding this comment.
I don't think we want to move these to the init.
Some new features like this actually relies on the fact that our optimizers lazily initialize their buffers.
torch/optim/functional.py
Outdated
| bias_correction2 = 1 - beta2 ** step | ||
|
|
||
| if weight_decay != 0: | ||
| grad = grad.add(p, alpha=weight_decay) |
There was a problem hiding this comment.
p or param? it looks like this line doesn't fall under any current tests?
There was a problem hiding this comment.
seems like the test_adam in test_optim didn't cover the weight_decay case, added a test case to cover that.
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
Codecov Report
@@ Coverage Diff @@
## gh/wanchaol/131/base #44791 +/- ##
========================================================
+ Coverage 68.11% 68.13% +0.01%
========================================================
Files 394 394
Lines 50957 50976 +19
========================================================
+ Hits 34710 34730 +20
+ Misses 16247 16246 -1
Continue to review full report at Codecov.
|
Stack from ghstack:
Differential Revision: D23935257