Looking for a simpler way of torch.tensor modification. Probably there is a way to apply Unfold to initial tensor directly.
input:
tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]])
output:
tensor([[0., 1., 3., 4.], [1., 2., 4., 5.], [3., 4., 6., 7.], [4., 5., 7., 8.]])
possible solution:
import torch t = torch.linspace(0., 8., steps=9) t1 = t.reshape(3,3) # starting point t2 = torch.flatten(t1) t3 = t2.reshape(1, 1, 1, -1) # unfold works with 4D only unfold = torch.nn.Unfold(kernel_size=(1, 5), dilation=1) t4 = unfold(t3) indices = torch.tensor([0, 1, 3, 4]) # deleting 3d (or middle) row and 3d (middle) column t5 = torch.index_select(torch.index_select(t4.squeeze(), 0, indices), 1, indices) t5
Advertisement
Answer
You can use unfold
, but in a simpler manner:
import torch import torch.nn.functional as nnf t1 = torch.arange(9.).reshape(3,3) # initial tensor out = nnf.unfold(t1[None, None, ...], kernel_size=2, padding=0) # that's it. done.