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