| |
| import torch |
| import torch.distributed as dist |
| import torch.nn as nn |
| from torch._utils import (_flatten_dense_tensors, _take_tensors, |
| _unflatten_dense_tensors) |
|
|
| from annotator.uniformer.mmcv.utils import TORCH_VERSION, digit_version |
| from .registry import MODULE_WRAPPERS |
| from .scatter_gather import scatter_kwargs |
|
|
|
|
| @MODULE_WRAPPERS.register_module() |
| class MMDistributedDataParallel(nn.Module): |
|
|
| def __init__(self, |
| module, |
| dim=0, |
| broadcast_buffers=True, |
| bucket_cap_mb=25): |
| super(MMDistributedDataParallel, self).__init__() |
| self.module = module |
| self.dim = dim |
| self.broadcast_buffers = broadcast_buffers |
|
|
| self.broadcast_bucket_size = bucket_cap_mb * 1024 * 1024 |
| self._sync_params() |
|
|
| def _dist_broadcast_coalesced(self, tensors, buffer_size): |
| for tensors in _take_tensors(tensors, buffer_size): |
| flat_tensors = _flatten_dense_tensors(tensors) |
| dist.broadcast(flat_tensors, 0) |
| for tensor, synced in zip( |
| tensors, _unflatten_dense_tensors(flat_tensors, tensors)): |
| tensor.copy_(synced) |
|
|
| def _sync_params(self): |
| module_states = list(self.module.state_dict().values()) |
| if len(module_states) > 0: |
| self._dist_broadcast_coalesced(module_states, |
| self.broadcast_bucket_size) |
| if self.broadcast_buffers: |
| if (TORCH_VERSION != 'parrots' |
| and digit_version(TORCH_VERSION) < digit_version('1.0')): |
| buffers = [b.data for b in self.module._all_buffers()] |
| else: |
| buffers = [b.data for b in self.module.buffers()] |
| if len(buffers) > 0: |
| self._dist_broadcast_coalesced(buffers, |
| self.broadcast_bucket_size) |
|
|
| def scatter(self, inputs, kwargs, device_ids): |
| return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim) |
|
|
| def forward(self, *inputs, **kwargs): |
| inputs, kwargs = self.scatter(inputs, kwargs, |
| [torch.cuda.current_device()]) |
| return self.module(*inputs[0], **kwargs[0]) |
|
|
| def train_step(self, *inputs, **kwargs): |
| inputs, kwargs = self.scatter(inputs, kwargs, |
| [torch.cuda.current_device()]) |
| output = self.module.train_step(*inputs[0], **kwargs[0]) |
| return output |
|
|
| def val_step(self, *inputs, **kwargs): |
| inputs, kwargs = self.scatter(inputs, kwargs, |
| [torch.cuda.current_device()]) |
| output = self.module.val_step(*inputs[0], **kwargs[0]) |
| return output |
|
|
