可以torch.cat
用来连接张量并获得所需的张量。这是一个完整的例子:
# input tensor
In [98]: at = torch.arange(1, 9).reshape(-1, 2).float()
In [99]: at
Out[99]:
tensor([[1., 2.],
[3., 4.],
[5., 6.],
[7., 8.]])
# columns to be padded
In [100]: col_zeros = torch.zeros(at.shape[0]).reshape(-1, 1)
In [101]: col_zeros
Out[101]:
tensor([[0.],
[0.],
[0.],
[0.]])
# rows to be padded
In [102]: row_zeros = torch.zeros(at.shape[1]+2).reshape(1, -1)
In [103]: row_zeros
Out[103]: tensor([[0., 0., 0., 0.]])
让我们先填充列:
# the order in the list of tensors matter.
# since we want a zero column on both sides, we place the input tensor in the middle
# and pad the `col_zeros` on both sides (i.e. along dimension=1)
In [104]: col_padded = torch.cat([col_zeros, at, col_zeros], dim=1)
In [105]: col_padded
Out[105]:
tensor([[0., 1., 2., 0.],
[0., 3., 4., 0.],
[0., 5., 6., 0.],
[0., 7., 8., 0.]])
接下来,让我们填充行:
# here we pad the `row_zeros` on the upper and lower sides (i.e. along dimension=0)
# placing the already `col_padded` tensor in the middle of the list of tensors
In [106]: final_padded = torch.cat([row_zeros, col_padded, row_zeros], dim=0)
In [107]: final_padded
Out[107]:
tensor([[0., 0., 0., 0.],
[0., 1., 2., 0.],
[0., 3., 4., 0.],
[0., 5., 6., 0.],
[0., 7., 8., 0.],
[0., 0., 0., 0.]])