Forcefully add dimension to NumPy array

Question

I have a four-dimensional data in array trainAll of shape N × H × W × 3. I need to separate it so I did

X_train = trainAll[:,:,:,1]
Y_train = trainAll[:,:,:,1:3]

As expected, Y_train.shape was N × H × W × 2.

But X_train.shape is N × H × W because the last dimension has just size 1. But neural network need four dimensional array, so it should look like

N × H × W × 1

The amazing thing is, if I do trainAll[:,:,:,2:3] then I get N*H*W*1 but I want the first dimension separated, not the last.

Honestly, I was unable to google because I did not know what to ask. So can any one help me out, so that I can not only separate first dimension but also shape is N × H × W × 1 instead of N × H × W ?

Answer 1

I have found following and it works much better: tf.expand_dims tensorflow docs, to reduce dimension instead use: tf.squeeze() here tf refers to tensorflow

Answer 2

Just try to add a new axis as the desired dimension. (Here, as the fourth dimension).

X_train = trainAll[:, :, :, 0]
X_train = X_train[:, :, :, np.newaxis]
# now, X_train.shape will be N * H * W * 1

The reason why you don't get them at the first place when you slice them is because slice hands the result as (n,) when using a single index and you make it (n, 1) by adding a new axis.

Answer 3

I was able to figure it out but still do not know if my answer is right. I wanted to know the python way to do it and What is happening when shape is shifted to N*H*W instead of N*H*W*1

Solution: trainAll[:,:,:,0:1] so instead of trainAll[:,:,:,1] picking it up just slice it