Search code examples
numpyreshapedimensions

NumPy: Can't add Dimension using .reshape()


As part of a Conv2D task I've been set, I have a training set that is shape (60000, 28, 28), this consists of number of training images, x axis and y axis dimensions respectively. I'm asked in the project I'm working on to add an additional dimension that would result in the shape being (60000, 28, 28, 1). One image should also return the shape (28, 28, 1) as a result.

I have tried images = training_images.reshape(60000, 28, 28, 1) however this only seems to return the shape as (60000, 28, 28) and does not return the extra dimension.


Solution

  • You can use expand_dims to add a dimension of size 1 in any of the array axis.

    import numpy as np
    
    a = np.zeros((400, 28, 28))
    b = np.expand_dims(a, axis=-1)
    b.shape
    
    > (400, 28, 28, 1)
    

    Or, if you need a shorter way (to use for example in an expression)

    a = np.zeros((400, 28, 28))
    a[..., np.newaxis].shape
    
    > (400, 28, 28, 1)
    

    where ... means every dimension of the array (it's like writing a[:, :, :, np.newaxis]).

    If you need to use reshape:

    a = np.zeros((400, 28, 28))
    a.reshape(*a.shape, 1).shape
    
    > (400, 28, 28, 1)
    

    *a.shape unpacks the shape tuple and passes each element as an argument of the reshape function (exactly as if you wrote a.reshape(400, 28, 28, 1))