PyTorch中为ImageFolder数据集随机分配标签的正确方法

问题背景 #

使用PyTorch的ImageFolder加载图像数据集时，默认会按照文件夹顺序分配标签（如class_1对应0，class_2对应1，以此类推）。若想为每张图片随机分配标签，直接修改dataset.targets的方式可能导致训练出错，以下是具体场景和问题：

原数据集加载代码 #

dataset = datasets.ImageFolder('...', transform=transform)
loader = DataLoader(dataset, batch_size=args.batchsize)

数据集文件夹结构 #

dataset/
  class_1/
  class_2/
  class_3/

尝试的错误修改方式 #

new_labels = [random.randint(0, 3) for i in range(len(dataset.targets))]
dataset.targets = new_labels

问题分析与解决方案 #

直接替换dataset.targets的问题在于：ImageFolder的标签系统与classes、class_to_idx属性绑定，仅修改targets会导致标签范围与数据集元信息不匹配（比如生成0-3的标签，但原classes只有3个类别），进而引发训练时的索引越界或损失计算错误。

推荐方法：自定义包装Dataset #

通过包装原ImageFolder，重写__getitem__方法来返回随机标签，这是最稳妥的方式：

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

class RandomLabelDataset(Dataset):
    def __init__(self, base_dataset, num_classes):
        self.base_dataset = base_dataset
        self.num_classes = num_classes
    
    def __len__(self):
        return len(self.base_dataset)
    
    def __getitem__(self, idx):
        img, _ = self.base_dataset[idx]
        # 生成0到num_classes-1范围内的随机标签
        random_label = random.randint(0, self.num_classes - 1)
        return img, random_label

# 加载原数据集
dataset = datasets.ImageFolder('...', transform=transform)
# 包装为随机标签数据集（示例设置4类标签）
random_dataset = RandomLabelDataset(dataset, num_classes=4)
loader = DataLoader(random_dataset, batch_size=args.batchsize)

可选方法：同步修改Dataset元信息 #

如果坚持直接修改原ImageFolder，需要同步更新classes和class_to_idx，确保标签与元信息匹配：

import random

dataset = datasets.ImageFolder('...', transform=transform)
# 生成0-3的随机标签
new_labels = [random.randint(0, 3) for _ in range(len(dataset.targets))]
dataset.targets = new_labels
# 更新类别元信息，对应新的标签范围
dataset.classes = ['class_0', 'class_1', 'class_2', 'class_3']
dataset.class_to_idx = {cls: idx for idx, cls in enumerate(dataset.classes)}

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