How to pad with zeros a tensor along some axis (Python)

Question

I would like to pad a numpy tensor with 0 along the chosen axis.
For instance, I have tensor r with shape (4,3,2) but I am only interested in padding only the last two axis (that is, pad only the matrix). Is it possible to do it with the one-line python code?

Asked By: m4linka

||

Source

Answer 1

You can use 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.]]]

Answered By: ali_m

Answer 2

This function would pad at the end of certain axis.
If you wish to pad both side, just modify it.

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)

example:

>>> 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)

Answered By: cswu

Answer 3

If you want a one-liner go with the solution proposed by @ali_m, but with that you should calculate and set by yourself the number of rows and columns with which to pad the input array. With the following function you can specify directly the target shape and it will return the array symmetrically padded with the value specified:

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)

Example

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

Answered By: Aelius

How to pad with zeros a tensor along some axis (Python)

Question:

Answers:

Example