Is there a way to use torch.cuda.comm.gather() with different tensor size like reduce_add_coalesced()?

cbats · November 14, 2017, 10:31am

torch.cuda.comm.reduce_add_coalesced() can handle a list of tensors with different size, but gather() can not.

pytorch/pytorch/blob/master/torch/cuda/comm.py#L196


outputs = []
for device, chunk, stream in zip(devices, chunks, streams):
    with torch.cuda.device(device), torch.cuda.stream(stream):
        outputs.append(chunk.cuda(device, async=True))
return tuple(outputs)




def gather(tensors, dim=0, destination=None):
"""Gathers tensors from multiple GPUs.


Tensor sizes in all dimension different than ``dim`` have to match.


Arguments:
    tensors (Iterable[Tensor]): iterable of tensors to gather.
    dim (int): a dimension along which the tensors will be concatenated.
    destination (int, optional): output device (-1 means CPU, default:
        current device)


Returns:
    A tensor located on ``destination`` device, that is a result of
    concatenating ``tensors`` along ``dim``.

Is there a way to make gather supporting different tensor size?

Thanks

alexis-jacq · November 14, 2017, 2:02pm

I don’t know if it would solve your problem, but you can use gather with different tensor sizes by repeating tensors along added dimensions:

values = torch.rand(3,7)
indices = (7*torch.rand(5,6)).long() # indices for values's 2nd dimension
repeated_values = values.view(3,7,1,1).expand(3,7,5,6)
repeated_indices = indices.view(1,1,5,6).expand(3,1,5,6) # note the 1 on the dimension for gathering
gathered = repeated_values.gather(1, repeated_indices).squeeze() # tensor of size 3,5,6

cbats · November 15, 2017, 7:13am

Sorry for my unclear description. I think you’re focus on the torch.gather method. And torch.cuda.comm.gather method is what I’m using which can communicate data between multiple GPUs and gather data to specific device.

Thanks

http://pytorch.org/docs/master/cuda.html?highlight=gather#torch.cuda.comm.gather