| 47 | |
| 48 | |
| 49 | class DistModule(torch.nn.Module): |
| 50 | def __init__(self, module, sync=False): |
| 51 | super(DistModule, self).__init__() |
| 52 | self.module = module |
| 53 | self.broadcast_params() |
| 54 | |
| 55 | self.sync = sync |
| 56 | if not sync: |
| 57 | self._grad_accs = [] |
| 58 | self._register_hooks() |
| 59 | |
| 60 | def forward(self, *inputs, **kwargs): |
| 61 | return self.module(*inputs, **kwargs) |
| 62 | |
| 63 | def _register_hooks(self): |
| 64 | for i, (name, p) in enumerate(self.named_parameters()): |
| 65 | if p.requires_grad: |
| 66 | p_tmp = p.expand_as(p) |
| 67 | grad_acc = p_tmp.grad_fn.next_functions[0][0] |
| 68 | grad_acc.register_hook(self._make_hook(name, p, i)) |
| 69 | self._grad_accs.append(grad_acc) |
| 70 | |
| 71 | def _make_hook(self, name, p, i): |
| 72 | def hook(*ignore): |
| 73 | link.allreduce_async(p.grad.data) |
| 74 | return hook |
| 75 | |
| 76 | def sync_gradients(self): |
| 77 | """ average gradients """ |
| 78 | if self.sync and link.get_world_size() > 1: |
| 79 | for name, param in self.module.named_parameters(): |
| 80 | if param.requires_grad and param.grad is not None: |
| 81 | link.allreduce(param.grad.data) |
| 82 | else: |
| 83 | link.synchronize() |
| 84 | |
| 85 | def broadcast_params(self): |
| 86 | """ broadcast model parameters """ |
| 87 | for name, param in self.module.state_dict().items(): |
| 88 | link.broadcast(param, 0) |
| 89 | |
| 90 | |
| 91 | def _serialize_to_tensor(data, group=None): |
no outgoing calls
no test coverage detected