Pytorch, эффективный способ расширить тензор по его первому и последнему элементу

Question

У меня есть тензор в pytorch.Я хочу расширить его на определенное измерение от начала и до конца на k позиций с первым и последним элементами этого измерения соответственно.
Скажем, у меня есть тензор с данными [[0, 0, 0], [1, 1, 1], [2, 2, 2]].Операция extend(dim, k) изменит ее следующим образом:

extend(0, 1): [ [0, 0, 0], [0, 0, 0] , [0, 0, 0], [1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2] ]

extend(1, 1):[ 0 , 0, 0, 0, 0 ], [ 1 , 1, 1, 1, 1 ], [ 2 , 2, 2, 2, 2 ]]

Какой эффективный способ сделать это (в соответствии с tensor.requires_grad=true)

Answer 1

Вы ищете torch.nn.functional.pad, с mode='replicate'.
Однако есть две вещи, на которые нужно обратить внимание, чтобы заставить это работать:
1. pad не работает с 2D тензорами.Таким образом, вам нужно добавить ведущие одноэлементные измерения до pad и squeeze, а затем их.
2. Порядок значений пэдов pad предполагает напротив против затемнения.

import torch
from torch.nn inport functional
x = torch.tensor([[0, 0, 0],[1, 1, 1], [2, 2, 2]], dtype=torch.float)
# expand along dim=0 by k=2 
f.pad(x[None,None,...], (0,0, 2, 2), mode='replicate').squeeze()

Out[]:
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [1., 1., 1.],
        [2., 2., 2.],
        [2., 2., 2.],
        [2., 2., 2.]])

# expand along dim=1 by k=2
f.pad(x[None,None,...], (2, 2, 0 , 0), mode='replicate').squeeze()

Out[]:
tensor([[0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2., 2., 2.]])