Search code examples
pytorch

Stretch the values of a pytorch tensor

I wish to 'stretch' the last two dimensions of a pytorch tensor to increase the spatial resolution of a (batch, channels, y, x) tensor.

Minimal example (I need 'new_function')

a = torch.tensor([[1, 2], [3, 4]])
b = new_function(a, (2, 3))
print(b)
tensor([[1, 1, 1, 2, 2, 2],
        [1, 1, 1, 2, 2, 2],
        [3, 3, 3, 4, 4, 4],
        [3, 3, 3, 4, 4, 4]])

One way of doing this (for the real problem):

a = torch.ones((2, 256, 2, 2)) # my original data.
b = torch.zeros((2, 256, 80, 96)) # The output I need
b[:, :, :40, :48] = a[:, :, 0, 0]
b[:, :, 40:, :48] = a[:, :, 1, 0]
b[:, :, :40, 48:] = a[:, :, 0, 1]
b[:, :, 40:, 48:] = a[:, :, 1, 1]
Solution

  • Use torch.nn.functional.interpolate (thanks to Shai)

    torch.nn.functional.interpolate(input_tensor.float(), size=(4, 6))

    My original idea was to use a variety of view and repeat methods:

    def stretch(e, sdims):
    od = e.shape
    return e.view(od[0], od[1], -1, 1).repeat(1, 1, 1, sdims[-1]).view(od[0], od[1], od[2], -1).repeat(1, 1, 1, sdims[-2]).view(od[0], od[1], od[2] * sdims[0], od[3] * sdims[1])

    torch.Size([2, 2, 4, 6])