| 58 | self.module.train(mode) |
| 59 | |
| 60 | def reduce_gradients(self, task_specific=False): |
| 61 | if self.sync: |
| 62 | if not task_specific: |
| 63 | if self.task_grp is not None or self.share_backbone_group is not None \ |
| 64 | or self.share_neck_group is not None or self.share_decoder_group is not None: |
| 65 | for name, param in self.named_parameters(): |
| 66 | if param.grad is None: param.grad = param.data * 0 |
| 67 | if param.task_specific and param.requires_grad: |
| 68 | allreduce(param.grad.data, group=self.task_grp) |
| 69 | elif param.backbone_specific and param.requires_grad: |
| 70 | allreduce(param.grad.data, group=self.share_backbone_group) |
| 71 | elif param.neck_specific and param.requires_grad: |
| 72 | allreduce(param.grad.data, group=self.share_neck_group) |
| 73 | elif param.decoder_specific and param.requires_grad: |
| 74 | allreduce(param.grad.data, group=self.share_decoder_group) |
| 75 | elif param.requires_grad: |
| 76 | allreduce(param.grad.data) |
| 77 | else: |
| 78 | for param in self.parameters(): |
| 79 | if param.requires_grad and param.grad is not None: |
| 80 | allreduce(param.grad.data) |
| 81 | else: |
| 82 | for name, param in self.named_parameters(): |
| 83 | if param.requires_grad and param.grad is not None: |
| 84 | allreduce(param.grad.data, group=self.task_grp) |
| 85 | |
| 86 | |
| 87 | def broadcast_params_multitask(model, task_grp, share_backbone_group, share_neck_group, share_decoder_group, ignore): |