| | |
| | 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 |
| |
|
