Implement the AMP training logic.
(self)
| 312 | self.grad_scaler = grad_scaler |
| 313 | |
| 314 | def run_step(self): |
| 315 | """ |
| 316 | Implement the AMP training logic. |
| 317 | """ |
| 318 | assert self.model.training, "[AMPTrainer] model was changed to eval mode!" |
| 319 | assert torch.cuda.is_available(), "[AMPTrainer] CUDA is required for AMP training!" |
| 320 | from torch.cuda.amp import autocast |
| 321 | |
| 322 | start = time.perf_counter() |
| 323 | data = next(self._data_loader_iter) |
| 324 | data_time = time.perf_counter() - start |
| 325 | |
| 326 | with autocast(): |
| 327 | loss_dict = self.model(data) |
| 328 | losses = sum(loss_dict.values()) |
| 329 | |
| 330 | self.optimizer.zero_grad() |
| 331 | self.grad_scaler.scale(losses).backward() |
| 332 | |
| 333 | self._write_metrics(loss_dict, data_time) |
| 334 | |
| 335 | self.grad_scaler.step(self.optimizer) |
| 336 | self.grad_scaler.update() |
nothing calls this directly
no test coverage detected