python - 如何在某个轴上用零填充张量（Python）

Question

我想沿所选轴用 0 填充一个 numpy 张量。例如，我有r带形状的张量，(4,3,2)但我只对仅填充最后两个轴感兴趣（即仅填充矩阵）。是否可以使用一行 python 代码来做到这一点？

Answer 1

您可以使用np.pad()：

a = np.ones((4, 3, 2))

# npad is a tuple of (n_before, n_after) for each dimension
npad = ((0, 0), (1, 2), (2, 1))
b = np.pad(a, pad_width=npad, mode='constant', constant_values=0)

print(b.shape)
# (4, 6, 5)

print(b)
# [[[ 0.  0.  0.  0.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  0.  0.  0.]
#   [ 0.  0.  0.  0.  0.]]

#  [[ 0.  0.  0.  0.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  0.  0.  0.]
#   [ 0.  0.  0.  0.  0.]]

#  [[ 0.  0.  0.  0.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  0.  0.  0.]
#   [ 0.  0.  0.  0.  0.]]

#  [[ 0.  0.  0.  0.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  1.  1.  0.]
#   [ 0.  0.  0.  0.  0.]
#   [ 0.  0.  0.  0.  0.]]]

Answer 2

此函数将在某个轴的末端填充。
如果你想填充两边，只需修改它。

def pad_along_axis(array: np.ndarray, target_length: int, axis: int = 0):

    pad_size = target_length - array.shape[axis]

    if pad_size <= 0:
        return array

    npad = [(0, 0)] * array.ndim
    npad[axis] = (0, pad_size)

    return np.pad(array, pad_width=npad, mode='constant', constant_values=0)

例子：

>>> a = np.identity(5)
>>> b = pad_along_axis(a, 7, axis=1)
>>> print(a, a.shape)
[[1. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0.]
 [0. 0. 1. 0. 0.]
 [0. 0. 0. 1. 0.]
 [0. 0. 0. 0. 1.]] (5, 5)

>>> print(b, b.shape)
[[1. 0. 0. 0. 0. 0. 0.]
 [0. 1. 0. 0. 0. 0. 0.]
 [0. 0. 1. 0. 0. 0. 0.]
 [0. 0. 0. 1. 0. 0. 0.]
 [0. 0. 0. 0. 1. 0. 0.]] (5, 7)