[Python] Error When Using Batch Size Greater Than 1 in PyTorch

I'm building a neural network to predict how an image will be partitioned during compression using VVC (Versatile Video Coding). The model takes a single Y-frame from a YUV420 image as input and uses a CSV file containing the ground truth block positions and sizes for training.

Input and Ground Truth

• 
  Input: A 1-frame YUV420 10-bit image.
  
  
• 
  Ground Truth: A CSV file with block positions, sizes, and additional partitioning flags.
  

example(388016_320x480_37.yuv) enter image description here

example(388016_320x480_37.csv) enter image description here

Problem Description:


I implemented train.py and dataset.py, but I'm encountering an error when setting batch_size > 1 in the DataLoader. With a batch size of 1, the model works correctly, but increasing the batch size leads to runtime errors.

Code Summary:


Below is a simplified version of my custom_collate_fn and DataLoader setup:

def custom_collate_fn(batch):
frames = [item[0] for item in batch] # Y-frame tensors
blocks = [item[1] for item in batch] # Block information
frames = torch.stack(frames, dim=0) # Stacking frames along batch dimension
return frames, blocks

dataloader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=custom_collate_fn
)

Observations:

• 
  When batch_size = 1, the blocks_batch in the training loop is a list containing a single set of block data.
  
  
• 
  With batch_size > 1, it becomes a list of lists, causing errors when indexing.
  

for i, (frame, blocks_batch) in enumerate(dataloader):
frame = frame.to(device) # Shape: [batch_size, 1, H, W]
blocks = blocks_batch[0] # Works with batch_size=1 but fails with larger sizes

My Assumption:


It seems the issue arises from handling blocks_batch when batch_size > 1. The nested list structure makes it difficult to handle multiple batches.

Request:

• 
  How can I adjust the custom_collate_fn or the training loop to handle batch sizes greater than 1 effectively?
  
  
• 
  If there's a better approach to batch-wise handling of variable-length block data, I'd appreciate any advice.
  

Thank you for your help!

File "C:\Users\Administrator\Documents\VVC_fast\test4\train.py", line 91, in <module>
loss1 = criterion(out_split, target_split)
File "C:\ProgramData\miniconda3\envs\newenv3\lib\site-packages\torch\nn\modules\module.py", line 1518, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "C:\ProgramData\miniconda3\envs\newenv3\lib\site-packages\torch\nn\modules\module.py", line 1527, in _call_impl
return forward_call(*args, **kwargs)
File "C:\ProgramData\miniconda3\envs\newenv3\lib\site-packages\torch\nn\modules\loss.py", line 725, in forward
return F.binary_cross_entropy_with_logits(input, target,
File "C:\ProgramData\miniconda3\envs\newenv3\lib\site-packages\torch\nn\functional.py", line 3193, in binary_cross_entropy_with_logits
raise ValueError(f"Target size ({target.size()}) must be the same as input size ({input.size()})")
ValueError: Target size (torch.Size([1, 1])) must be the same as input size (torch.Size([2, 1]))

Continue reading...