| 24 | allgather = dist.all_gather |
| 25 | |
| 26 | class DistModule(torch.nn.Module): |
| 27 | def __init__(self, module, sync=False, task_grp=None, share_backbone_group=None, \ |
| 28 | share_neck_group=None, share_decoder_group=None, ignore_bcast=None, \ |
| 29 | task_weight=None, task_size=None): |
| 30 | super(DistModule, self).__init__() |
| 31 | self.module = module |
| 32 | self.sync = sync |
| 33 | self.task_grp = task_grp |
| 34 | self.share_backbone_group = share_backbone_group |
| 35 | self.share_neck_group = share_neck_group |
| 36 | self.share_decoder_group = share_decoder_group |
| 37 | self.task_weight = task_weight |
| 38 | self.task_size = task_size |
| 39 | |
| 40 | if not hasattr(torch.nn.Module, 'named_buffers'): |
| 41 | printlog('registering named_buffers for nn.Module at DistModule') |
| 42 | torch.nn.Module.named_buffers = named_buffers |
| 43 | |
| 44 | broadcast_params_multitask(self, self.task_grp, self.share_backbone_group, \ |
| 45 | self.share_neck_group, self.share_decoder_group, ignore_bcast) |
| 46 | |
| 47 | assert sync, "Currently, only sync model is supported!" |
| 48 | if not sync: |
| 49 | self._grad_accs = {} |
| 50 | self._reduce_hooks = {} |
| 51 | self._register_hooks() |
| 52 | |
| 53 | def forward(self, *inputs, **kwargs): |
| 54 | return self.module(*inputs, **kwargs) |
| 55 | |
| 56 | def train(self, mode=True): |
| 57 | super(DistModule, self).train(mode) |
| 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: |
no outgoing calls