How to iterate over the first index of 4D matrix every nth element?-CodePudding

Assume there is a 4D matrix a=(20,3,100,100). I want to divide the first index of a which is 20 into 4 batches within a for loop and then append each iteration into output. I don't want to use reshape since each batch will be fed into a network.

The size of the output should be (4,5,3,100,100)

Here is my try:

output=[]
a=np.random.randint(0,100,size=(20,3,600,600))
for i in a[1:20:5,:,:,:] :  
            cnn_in=a[i,:,:,:]
            output.append(cnn_in)

CodePudding user response：

something like this?

for i in range(0,20,5):
     cnn_in=a[i:i 5]
     output.append(cnn_in)

or simpler with np.split: output=np.split(a,4)