Implement the standard training logic described above.
(self)
| 211 | self.optimizer = optimizer |
| 212 | |
| 213 | def run_step(self): |
| 214 | """ |
| 215 | Implement the standard training logic described above. |
| 216 | """ |
| 217 | assert self.model.training, "[SimpleTrainer] model was changed to eval mode!" |
| 218 | start = time.perf_counter() |
| 219 | """ |
| 220 | If you want to do something with the data, you can wrap the dataloader. |
| 221 | """ |
| 222 | data = next(self._data_loader_iter) |
| 223 | data_time = time.perf_counter() - start |
| 224 | |
| 225 | """ |
| 226 | If you want to do something with the losses, you can wrap the model. |
| 227 | """ |
| 228 | loss_dict = self.model(data) |
| 229 | losses = sum(loss_dict.values()) |
| 230 | |
| 231 | """ |
| 232 | If you need to accumulate gradients or do something similar, you can |
| 233 | wrap the optimizer with your custom `zero_grad()` method. |
| 234 | """ |
| 235 | self.optimizer.zero_grad() |
| 236 | losses.backward() |
| 237 | |
| 238 | self._write_metrics(loss_dict, data_time) |
| 239 | |
| 240 | """ |
| 241 | If you need gradient clipping/scaling or other processing, you can |
| 242 | wrap the optimizer with your custom `step()` method. But it is |
| 243 | suboptimal as explained in https://arxiv.org/abs/2006.15704 Sec 3.2.4 |
| 244 | """ |
| 245 | self.optimizer.step() |
| 246 | |
| 247 | def _write_metrics(self, loss_dict: Dict[str, torch.Tensor], data_time: float): |
| 248 | """ |
nothing calls this directly
no test coverage detected