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

Question

我想在numpy张量上沿着指定的轴用0进行填充。 例如，我有形状为(4,3,2)的张量r，但我只想填充最后两个轴（即仅填充矩阵）。是否可以使用一行python代码实现？

- m4linka

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

- ali_m

---

请注意，(n_before, n_after) 是指行数/列数；因此，上面的 (1,2)，对于第二维来说，表示一个在其上方（之前）和两个在其下方（之后）的行。同样地，(2,1) 表示在其左侧向右移动 2 列 0 填充，并通过添加 1 列零进行左移填充的 2 行。 - berto77

@berto77 谢谢你的评论，我很苦恼。 - Euler_Salter

Answer 2

这个函数会在特定轴的末尾进行填充。
如果你想要两边都填充，只需修改它。

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

    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)

例子：

>>> import numpy as np
>>> 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)

- cswu

---

len(array.shape) == array.ndim。如果 pad_size <= 0，你可以安全地 return。最后，npad = [(0, 0)] * array.ndim 更短 :) - bers

Answer 3

如果你想要一行代码，可以使用@ali_m提出的解决方案，但是你需要自己计算并设置要用来填充输入数组的行数和列数。使用下面的函数，你可以直接指定目标形状，它将返回以指定值对称填充的数组：

def symmetric_pad_array(input_array: np.ndarray, target_shape: tuple, pad_value: int) -> np.ndarray:

    for dim_in, dim_target in zip(input_array.shape, target_shape):
        if dim_target < dim_in:
            raise Exception("`target_shape` should be greater or equal than `input_array` shape for each axis.")

    pad_width = []
    for dim_in, dim_target  in zip(input_array.shape, target_shape):
        if (dim_in-dim_target)%2 == 0:
            pad_width.append((int(abs((dim_in-dim_target)/2)), int(abs((dim_in-dim_target)/2))))
        else:
            pad_width.append((int(abs((dim_in-dim_target)/2)), (int(abs((dim_in-dim_target)/2))+1)))
    
    return np.pad(input_array, pad_width, 'constant', constant_values=pad_value)

例子

>>> a = np.array(np.arange(0,27)).reshape(3,3,3)
>>> target_shape = (3,5,5)
>>> symmetric_pad_array(a, target_shape, pad_value=0)
array([[[ 0,  0,  0,  0,  0],
        [ 0,  0,  1,  2,  0],
        [ 0,  3,  4,  5,  0],
        [ 0,  6,  7,  8,  0],
        [ 0,  0,  0,  0,  0]],

       [[ 0,  0,  0,  0,  0],
        [ 0,  9, 10, 11,  0],
        [ 0, 12, 13, 14,  0],
        [ 0, 15, 16, 17,  0],
        [ 0,  0,  0,  0,  0]],

       [[ 0,  0,  0,  0,  0],
        [ 0, 18, 19, 20,  0],
        [ 0, 21, 22, 23,  0],
        [ 0, 24, 25, 26,  0],
        [ 0,  0,  0,  0,  0]]])

- Aelius