Iterate over 5d tensor


I want to iterate over a 5d tensor ‘images’, extract 4d tensors and do some operation on the 4d tensors which is an individual bag of (batch size, channels, height, width).So in a 5d tensor, 1st dimension is batch size, 2nd dimension is bag size which is number of images in a bag (lets say 20),3rd dimension is RGB, 4th and 5th dimension are height and is my for loop which I am using but this very there any other way to make this code faster? the for loop is very slow.

                new_output = []
                for i in range(args.bag_size):
                    individual_bag = images[:,i,:,:,:]
                    print('individual_bag shape',individual_bag.shape)
                    intermediate_output = model.get_intermediate_layers(individual_bag, n)
                    output =[x[:, 0] for x in intermediate_output], dim=-1)
                output = torch.mean(torch.stack(new_output),dim = 0)


I don’t know what each “individual bag” represents, but since you are indexing each separately I assume they don’t have any dependency between each other. If so, would it be possible to flatten dim1 into dim0 and just increase the batch size for a single forward pass?

Each individual bag represents lets say 20 images, so instead of passing each image, I am passing a bag of images.I am extracting features from the network for each bag in a batch, and doing this for each batch.The purpose of taking a bag of images is to get mean representation but I think I would need to do it for each batch separately.

I’m unsure why you would need to perform this operation for each bag separately. You could still reshape the output back before applying the mean operation.