Set weight decay and step optimizer.
(self)
| 130 | |
| 131 | # Pytorch optimizer methods |
| 132 | def step(self) -> None: |
| 133 | "Set weight decay and step optimizer." |
| 134 | # weight decay outside of optimizer step (AdamW) |
| 135 | if self.true_wd: |
| 136 | for lr, wd, pg1, pg2 in zip(self._lr, self._wd, self.opt.param_groups[::2], self.opt.param_groups[1::2]): |
| 137 | for p in pg1['params']: |
| 138 | # When some parameters are fixed: Shaoshuai Shi |
| 139 | if p.requires_grad is False: |
| 140 | continue |
| 141 | p.data.mul_(1 - wd * lr) |
| 142 | if self.bn_wd: |
| 143 | for p in pg2['params']: |
| 144 | # When some parameters are fixed: Shaoshuai Shi |
| 145 | if p.requires_grad is False: |
| 146 | continue |
| 147 | p.data.mul_(1 - wd * lr) |
| 148 | self.set_val('weight_decay', listify(0, self._wd)) |
| 149 | self.opt.step() |
| 150 | |
| 151 | def zero_grad(self) -> None: |
| 152 | "Clear optimizer gradients." |