Multi GPU training, memory usage in-balance

HANG_ZHANG · June 23, 2017, 4:43am

I found this:

pytorch/pytorch/blob/master/torch/nn/parallel/data_parallel.py#L47


    >>> output = net(input_var)
"""


# TODO: update notes/cuda.rst when this class handles 8+ GPUs well


def __init__(self, module, device_ids=None, output_device=None, dim=0):
    super(DataParallel, self).__init__()


    if not torch.cuda.is_available():
        self.module = module
        self.device_ids = []
        return


    if device_ids is None:
        device_ids = list(range(torch.cuda.device_count()))
    if output_device is None:
        output_device = device_ids[0]
    self.dim = dim
    self.module = module
    self.device_ids = device_ids
    self.output_device = output_device

It seems always gathers the output to the first GPU. Is this a temporary solution?