Растянуть значения тензора пироха

Question

Я бы хотел «растянуть» последние два измерения тензора пироха, чтобы увеличить пространственное разрешение тензора (партия, каналы, у, х).

Минимальный пример (мне нужна функция '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]])

Один из способов сделать это (для реальной проблемы):

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]

Answer 1

Использование torch.nn.functional.interpolate (спасибо Shai)

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

Моя оригинальная идея заключалась в том, чтобы использовать различные методы view и repeat:

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