datasets.Dataset.set_transform()未成功应用图片变换的问题

核心问题 #

你直接把PyTorch的Compose([Resize(), ToTensor()])传给了Hugging Face Dataset的set_transform/with_transform，但这两个方法要求变换函数必须接收完整的样本字典（而非单独的图像）。这种写法相当于把整个样本字典当成图像输入给Resize，完全没处理到image字段，自然变换不会生效，图像还是原始尺寸，导致DataLoader堆叠时报错。

解决步骤 #

1. 编写针对样本字典的自定义变换函数 #

把PyTorch的图像变换包装成一个能处理样本字典的函数，明确对image字段进行变换：

from torchvision.transforms import Compose, Resize, ToTensor

def apply_transforms(sample):
    # 定义针对单张图像的变换链，指定固定宽高确保尺寸统一
    img_transforms = Compose([
        Resize((224, 224)),
        ToTensor()
    ])
    # 对样本中的image字段应用变换
    sample["image"] = img_transforms(sample["image"])
    return sample

2. 给数据集绑定变换 #

用set_transform或with_transform绑定这个自定义函数：

# 方式1：直接修改原数据集
dataset.set_transform(apply_transforms)

# 方式2：创建新的变换后数据集（推荐，不修改原始数据）
transformed_dataset = dataset.with_transform(apply_transforms)

3. 额外检查：图像格式是否为PIL Image #

如果数据集加载后image字段是numpy数组（而非PIL Image），需要在变换链中加入ToPILImage()：

from torchvision.transforms import ToPILImage

def apply_transforms(sample):
    img_transforms = Compose([
        ToPILImage(),  # 先转成PIL图像才能使用Resize
        Resize((224, 224)),
        ToTensor()
    ])
    sample["image"] = img_transforms(sample["image"])
    return sample

4. 验证变换是否生效 #

可以取一个样本查看尺寸，确认变换已生效：

sample = transformed_dataset[0]
print(sample["image"].shape)  # 正常会输出 torch.Size([3, 224, 224])

为什么之前的写法无效？ #

Hugging Face Dataset的set_transform会把整个样本字典传递给你传入的函数。如果你直接传Compose([Resize(), ToTensor()])，这个函数会尝试把字典当成图像处理，完全不会修改image字段，所以原始图像尺寸没有任何变化，最终导致DataLoader批量堆叠时因尺寸不一致报错。

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