|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| """Encoding Data Parallel"""
|
| import threading
|
| import functools
|
| import torch
|
| from torch.autograd import Variable, Function
|
| import torch.cuda.comm as comm
|
| from torch.nn.parallel.data_parallel import DataParallel
|
| from torch.nn.parallel.parallel_apply import get_a_var
|
| from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
|
|
|
| torch_ver = torch.__version__[:3]
|
|
|
| __all__ = ['allreduce', 'DataParallelModel', 'DataParallelCriterion', 'patch_replication_callback']
|
|
|
| def allreduce(*inputs):
|
| """Cross GPU all reduce autograd operation for calculate mean and
|
| variance in SyncBN.
|
| """
|
| return AllReduce.apply(*inputs)
|
|
|
| class AllReduce(Function):
|
| @staticmethod
|
| def forward(ctx, num_inputs, *inputs):
|
| ctx.num_inputs = num_inputs
|
| ctx.target_gpus = [inputs[i].get_device() for i in range(0, len(inputs), num_inputs)]
|
| inputs = [inputs[i:i + num_inputs]
|
| for i in range(0, len(inputs), num_inputs)]
|
|
|
| inputs = sorted(inputs, key=lambda i: i[0].get_device())
|
| results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0])
|
| outputs = comm.broadcast_coalesced(results, ctx.target_gpus)
|
| return tuple([t for tensors in outputs for t in tensors])
|
|
|
| @staticmethod
|
| def backward(ctx, *inputs):
|
| inputs = [i.data for i in inputs]
|
| inputs = [inputs[i:i + ctx.num_inputs]
|
| for i in range(0, len(inputs), ctx.num_inputs)]
|
| results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0])
|
| outputs = comm.broadcast_coalesced(results, ctx.target_gpus)
|
| return (None,) + tuple([Variable(t) for tensors in outputs for t in tensors])
|
|
|
| class Reduce(Function):
|
| @staticmethod
|
| def forward(ctx, *inputs):
|
| ctx.target_gpus = [inputs[i].get_device() for i in range(len(inputs))]
|
| inputs = sorted(inputs, key=lambda i: i.get_device())
|
| return comm.reduce_add(inputs)
|
|
|
| @staticmethod
|
| def backward(ctx, gradOutput):
|
| return Broadcast.apply(ctx.target_gpus, gradOutput)
|
|
|
|
|
| class DataParallelModel(DataParallel):
|
| """Implements data parallelism at the module level.
|
|
|
| This container parallelizes the application of the given module by
|
| splitting the input across the specified devices by chunking in the
|
| batch dimension.
|
| In the forward pass, the module is replicated on each device,
|
| and each replica handles a portion of the input. During the backwards pass, gradients from each replica are summed into the original module.
|
| Note that the outputs are not gathered, please use compatible
|
| :class:`encoding.parallel.DataParallelCriterion`.
|
|
|
| The batch size should be larger than the number of GPUs used. It should
|
| also be an integer multiple of the number of GPUs so that each chunk is
|
| the same size (so that each GPU processes the same number of samples).
|
|
|
| Args:
|
| module: module to be parallelized
|
| device_ids: CUDA devices (default: all devices)
|
|
|
| Reference:
|
| Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi,
|
| Amit Agrawal. “Context Encoding for Semantic Segmentation.
|
| *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018*
|
|
|
| Example::
|
|
|
| >>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2])
|
| >>> y = net(x)
|
| """
|
| def gather(self, outputs, output_device):
|
| return outputs
|
|
|
| def replicate(self, module, device_ids):
|
| modules = super(DataParallelModel, self).replicate(module, device_ids)
|
| return modules
|
|
|
|
|
| class DataParallelCriterion(DataParallel):
|
| """
|
| Calculate loss in multiple-GPUs, which balance the memory usage for
|
| Semantic Segmentation.
|
|
|
| The targets are splitted across the specified devices by chunking in
|
| the batch dimension. Please use together with :class:`encoding.parallel.DataParallelModel`.
|
|
|
| Reference:
|
| Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi,
|
| Amit Agrawal. “Context Encoding for Semantic Segmentation.
|
| *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018*
|
|
|
| Example::
|
|
|
| >>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2])
|
| >>> criterion = encoding.nn.DataParallelCriterion(criterion, device_ids=[0, 1, 2])
|
| >>> y = net(x)
|
| >>> loss = criterion(y, target)
|
| """
|
| def forward(self, inputs, *targets, **kwargs):
|
|
|
|
|
| if not self.device_ids:
|
| return self.module(inputs, *targets, **kwargs)
|
| targets, kwargs = self.scatter(targets, kwargs, self.device_ids)
|
| if len(self.device_ids) == 1:
|
| return self.module(inputs, *targets[0], **kwargs[0])
|
| replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
|
| outputs = _criterion_parallel_apply(replicas, inputs, targets, kwargs)
|
| return Reduce.apply(*outputs) / len(outputs)
|
|
|
|
|
| def _criterion_parallel_apply(modules, inputs, targets, kwargs_tup=None, devices=None):
|
| assert len(modules) == len(inputs)
|
| assert len(targets) == len(inputs)
|
| if kwargs_tup:
|
| assert len(modules) == len(kwargs_tup)
|
| else:
|
| kwargs_tup = ({},) * len(modules)
|
| if devices is not None:
|
| assert len(modules) == len(devices)
|
| else:
|
| devices = [None] * len(modules)
|
|
|
| lock = threading.Lock()
|
| results = {}
|
| if torch_ver != "0.3":
|
| grad_enabled = torch.is_grad_enabled()
|
|
|
| def _worker(i, module, input, target, kwargs, device=None):
|
| if torch_ver != "0.3":
|
| torch.set_grad_enabled(grad_enabled)
|
| if device is None:
|
| device = get_a_var(input).get_device()
|
| try:
|
| if not isinstance(input, tuple):
|
| input = (input,)
|
| with torch.cuda.device(device):
|
| output = module(*(input + target), **kwargs)
|
| with lock:
|
| results[i] = output
|
| except Exception as e:
|
| with lock:
|
| results[i] = e
|
|
|
| if len(modules) > 1:
|
| threads = [threading.Thread(target=_worker,
|
| args=(i, module, input, target,
|
| kwargs, device),)
|
| for i, (module, input, target, kwargs, device) in
|
| enumerate(zip(modules, inputs, targets, kwargs_tup, devices))]
|
|
|
| for thread in threads:
|
| thread.start()
|
| for thread in threads:
|
| thread.join()
|
| else:
|
| _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])
|
|
|
| outputs = []
|
| for i in range(len(inputs)):
|
| output = results[i]
|
| if isinstance(output, Exception):
|
| raise output
|
| outputs.append(output)
|
| return outputs
|
|
|
