| 448 | torch.nn.utils.clip_grad.clip_grad_norm_(self.model.parameters(), max_norm=self.clip_grad) |
| 449 | |
| 450 | def update(self): |
| 451 | ginfo = self.ginfo |
| 452 | tmp = self.tmp |
| 453 | |
| 454 | # reduce |
| 455 | self.model.reduce_gradients() |
| 456 | |
| 457 | if self.clip_grad_backbone > 0: |
| 458 | tmp.backbone_grad_norm = torch.nn.utils.clip_grad.clip_grad_norm_(\ |
| 459 | self.model.module.backbone_module.parameters(), \ |
| 460 | max_norm=self.clip_grad_backbone*(ginfo.task_size**0.5)) |
| 461 | is_inf = np.isinf(tmp.backbone_grad_norm) |
| 462 | is_nan = np.isnan(tmp.backbone_grad_norm) |
| 463 | if ginfo.task_rank == 0 and (is_inf or is_nan): |
| 464 | self.logger.info('task{} {} backbone_grad_norm inf/nan {}/{}'.format(\ |
| 465 | ginfo.task_id, ginfo.task_name, is_inf, is_nan)) |
| 466 | |
| 467 | if self.clip_grad_neck > 0: |
| 468 | tmp.neck_grad_norm = torch.nn.utils.clip_grad.clip_grad_norm_(\ |
| 469 | self.model.module.neck_module.parameters(), \ |
| 470 | max_norm=self.clip_grad_neck*(self.C.world_size**0.5)) |
| 471 | is_inf = np.isinf(tmp.neck_grad_norm) |
| 472 | is_nan = np.isnan(tmp.neck_grad_norm) |
| 473 | if ginfo.task_rank == 0 and (is_inf or is_nan): |
| 474 | self.logger.info('task{} {} backbone_grad_norm inf/nan {}/{}'.format(\ |
| 475 | ginfo.task_id, ginfo.task_name, is_inf, is_nan)) |
| 476 | |
| 477 | if self.clip_grad_decoder > 0: |
| 478 | tmp.decoder_grad_norm = torch.nn.utils.clip_grad.clip_grad_norm_(\ |
| 479 | self.model.module.decoder_module.parameters(), \ |
| 480 | max_norm=self.clip_grad_decoder*(self.C.world_size**0.5)) |
| 481 | is_inf = np.isinf(tmp.decoder_grad_norm) |
| 482 | is_nan = np.isnan(tmp.decoder_grad_norm) |
| 483 | if ginfo.task_rank == 0 and (is_inf or is_nan): |
| 484 | self.logger.info('task{} {} backbone_grad_norm inf/nan {}/{}'.format(\ |
| 485 | ginfo.task_id, ginfo.task_name, is_inf, is_nan)) |
| 486 | |
| 487 | self.optimizer.step() |
| 488 | |
| 489 | def tb_logging(self): |
| 490 | tmp = self.tmp |