PyTorch加载PascalVOC数据集报错：RuntimeError: batch元素大小不一致

你遇到的RuntimeError: each element in list of batch should be of equal size不是图像尺寸的问题——你已经用Resize把图像统一成448x448了，问题出在标注数据上。

VOCDetection返回的是(图像, 标注字典)的组合，标注字典里的object字段（包含边界框、类别等）是每个样本长度不同的列表（比如有的图有3个物体，有的有5个）。DataLoader默认会尝试把这些不同长度的标注张量堆叠，自然会报尺寸不匹配的错。

下面给两种解决方案：

方案一：只加载图像（用于测试图像预处理） #

如果只是验证图像是否能正常批量加载，可以自定义一个只返回图像的数据集包装类：

from torchvision import datasets
from torchvision.transforms import v2, ToTensor
from torch.utils.data import DataLoader, Dataset

class VOCImageOnly(Dataset):
    def __init__(self, root, image_set, transform=None):
        self.voc_data = datasets.voc.VOCDetection(root=root, image_set=image_set, download=False)
        self.transform = transform
    
    def __len__(self):
        return len(self.voc_data)
    
    def __getitem__(self, idx):
        img, _ = self.voc_data[idx]
        if self.transform:
            img = self.transform(img)
        return img

# 初始化数据集和DataLoader
validation_data = VOCImageOnly(
    root='./DATA/',
    image_set="val",
    transform=v2.Compose([v2.Resize(size=(448, 448)), ToTensor()])
)

validation_dataloader = DataLoader(validation_data, batch_size=64)

# 测试加载
for X in validation_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    break

方案二：保留标注（用于YOLO训练） #

如果需要用到标注数据，必须自定义collate_fn函数，告诉DataLoader如何处理不同长度的标注：

from torchvision import datasets
from torchvision.transforms import v2, ToTensor
from torch.utils.data import DataLoader
import torch

def yolo_collate_fn(batch):
    images = []
    targets = []
    for img, target in batch:
        images.append(img)
        # 提取标注中的物体信息
        obj_list = target['annotation']['object']
        # 处理单物体的情况（此时obj_list是字典而非列表）
        if isinstance(obj_list, dict):
            obj_list = [obj_list]
        
        bbox_list = []
        label_list = []
        orig_width = float(target['annotation']['size']['width'])
        orig_height = float(target['annotation']['size']['height'])
        
        for obj in obj_list:
            # 读取原始边界框并缩放到448x448尺寸
            bbox = obj['bndbox']
            xmin = float(bbox['xmin']) / orig_width * 448
            ymin = float(bbox['ymin']) / orig_height * 448
            xmax = float(bbox['xmax']) / orig_width * 448
            ymax = float(bbox['ymax']) / orig_height * 448
            bbox_list.append([xmin, ymin, xmax, ymax])
            
            # 这里需要替换成实际的VOC类别到索引的映射，示例用0代替
            label_list.append(0)
        
        targets.append({
            'boxes': torch.tensor(bbox_list),
            'labels': torch.tensor(label_list)
        })
    
    # 堆叠图像成张量，标注保留为列表形式
    images = torch.stack(images)
    return images, targets

# 初始化数据集和DataLoader，指定collate_fn
validation_data = datasets.voc.VOCDetection(
    root='./DATA/',
    download=False,
    image_set="val",
    transform=v2.Compose([v2.Resize(size=(448, 448)), ToTensor()])
)

validation_dataloader = DataLoader(validation_data, batch_size=64, collate_fn=yolo_collate_fn)

# 测试加载
for X, y in validation_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"第一个样本的边界框数量: {y[0]['boxes'].shape[0]}")
    break

说明：collate_fn的作用是自定义批量数据的拼接逻辑，这里把图像统一堆叠成张量，标注则以列表形式保存每个样本的独立信息，避免了尺寸不匹配的问题。

内容的提问来源于stack exchange，提问作者Henrique Hott