(self, batch, phase='train')
| 228 | self.logger.log_info('Sample done, time: {:.2f}'.format(time.time() - tic)) |
| 229 | |
| 230 | def step(self, batch, phase='train'): |
| 231 | loss = {} |
| 232 | if self.debug == False: |
| 233 | for k, v in batch.items(): |
| 234 | if torch.is_tensor(v): |
| 235 | batch[k] = v.cuda() |
| 236 | else: |
| 237 | batch = batch[0].cuda() |
| 238 | for op_sc_n, op_sc in self.optimizer_and_scheduler.items(): |
| 239 | if phase == 'train': |
| 240 | # check if this optimizer and scheduler is valid in this iteration and epoch |
| 241 | if op_sc['start_iteration'] > self.last_iter: |
| 242 | continue |
| 243 | if op_sc['end_iteration'] > 0 and op_sc['end_iteration'] <= self.last_iter: |
| 244 | continue |
| 245 | if op_sc['start_epoch'] > self.last_epoch: |
| 246 | continue |
| 247 | if op_sc['end_epoch'] > 0 and op_sc['end_epoch'] <= self.last_epoch: |
| 248 | continue |
| 249 | |
| 250 | input = { |
| 251 | 'batch': batch, |
| 252 | 'return_loss': True, |
| 253 | 'step': self.last_iter, |
| 254 | } |
| 255 | if op_sc_n != 'none': |
| 256 | input['name'] = op_sc_n |
| 257 | |
| 258 | if phase == 'train': |
| 259 | if self.args.amp: |
| 260 | with autocast(): |
| 261 | output = self.model(**input) |
| 262 | else: |
| 263 | output = self.model(**input) |
| 264 | else: |
| 265 | with torch.no_grad(): |
| 266 | if self.args.amp: |
| 267 | with autocast(): |
| 268 | output = self.model(**input) |
| 269 | else: |
| 270 | output = self.model(**input) |
| 271 | |
| 272 | if phase == 'train': |
| 273 | if op_sc['optimizer']['step_iteration'] > 0 and (self.last_iter + 1) % op_sc['optimizer']['step_iteration'] == 0: |
| 274 | op_sc['optimizer']['module'].zero_grad() |
| 275 | if self.args.amp: |
| 276 | self.scaler.scale(output['loss']).backward() |
| 277 | if self.clip_grad_norm is not None: |
| 278 | self.clip_grad_norm(self.model.parameters()) |
| 279 | self.scaler.step(op_sc['optimizer']['module']) |
| 280 | self.scaler.update() |
| 281 | else: |
| 282 | output['loss'].backward() |
| 283 | if self.clip_grad_norm is not None: |
| 284 | self.clip_grad_norm(self.model.parameters()) |
| 285 | op_sc['optimizer']['module'].step() |
| 286 | |
| 287 | if 'scheduler' in op_sc: |
no test coverage detected