| 66 | |
| 67 | # Helper function to extract optimizer states from a torch.optim.Adam instance |
| 68 | def get_inputs(optim): |
| 69 | steps = [] |
| 70 | params = [] |
| 71 | grads = [] |
| 72 | exp_avgs = [] |
| 73 | exp_avg_sqs = [] |
| 74 | for group in optim.param_groups: |
| 75 | for p in group["params"]: |
| 76 | params.append(p) |
| 77 | grads.append(p.grad) |
| 78 | state = optim.state[p] |
| 79 | exp_avgs.append(state["exp_avg"]) |
| 80 | exp_avg_sqs.append(state["exp_avg_sq"]) |
| 81 | steps.append(state["step"]) |
| 82 | |
| 83 | return steps, params, exp_avgs, exp_avg_sqs |
| 84 | |
| 85 | |
| 86 | # Functions to update the different optimizer states |