Zero fp32 and fp16 parameter grads.
(self, set_grads_to_None=False)
| 214 | "load_state_dict().") |
| 215 | |
| 216 | def zero_grad(self, set_grads_to_None=False): |
| 217 | """ |
| 218 | Zero fp32 and fp16 parameter grads. |
| 219 | """ |
| 220 | # In principle, only the .grad attributes of the model params need to be |
| 221 | # zeroed, because gradients are copied into the FP32 master params. |
| 222 | # However, we zero all gradients owned by the optimizer, just to be safe: |
| 223 | for group in self.optimizer.param_groups: |
| 224 | for p in group['params']: |
| 225 | if set_grads_to_None: |
| 226 | p.grad = None |
| 227 | else: |
| 228 | if p.grad is not None: |
| 229 | p.grad.detach_() |
| 230 | p.grad.zero_() |
| 231 | |
| 232 | # Zero fp16 gradients owned by the model: |
| 233 | for fp16_group in self.fp16_groups: |
| 234 | for param in fp16_group: |
| 235 | if set_grads_to_None: |
| 236 | param.grad = None |
| 237 | else: |
| 238 | if param.grad is not None: |
| 239 | # as in torch.optim.optimizer.zero_grad() |
| 240 | param.grad.detach_() |
| 241 | param.grad.zero_() |
| 242 | |
| 243 | def _master_params_to_model_params(self): |
| 244 | for fp16_group, fp32_from_fp16_group in zip(self.fp16_groups, |