How to remove batch size in training loop to perform certain operations on a single image

1

I am trying to perform certain operations on a single image, while in a training loop.
In case of batch_size = 1 , it could be easily done by using torch.squeeze but I am unable to think of a way when I can do it for other batch sizes.
Below is the minimum code for representation -

def train(n_epochs, loaders, model, optimizer, criterion, use_cuda, save_path):
      for epoch in range(1, n_epochs+1):
             for batch_idx, (data, target) in enumerate(final_train_loader):
                        # Here the target shape would be B*H*W*N 

B: Batch size
H,W,N:Height,width,no. of channels

I want it in a form of H * W * N

2

If you want to perform some operations on single samples, you could use a for loop and iterate your batch:

for d, t in zip(data, target):
    # op on single sample
2 Likes