PyTorch图像数据集分层划分时遇TypeError：期望Tensor Image求助

错误根源 #

1. Transform顺序错误：transforms.Normalize仅能处理Tensor类型数据，但你把它放在了transforms.ToTensor()之前，此时输入是PIL图像，本身就会引发类型不匹配问题。
2. 冗余的二次转换：Subset会继承原始数据集的transform逻辑，输出的已经是Tensor格式图像，你后续用PILToTensor()去转换Tensor，自然触发img should be Tensor Image错误——因为PILToTensor()只接受PIL图像作为输入。

修正方案 #

步骤1：调整Transform执行顺序 #

将transforms.ToTensor()移到transforms.Normalize之前，确保Normalize处理的是Tensor数据：

imageSize = 224
train_transforms = transforms.Compose([
    transforms.Resize((imageSize, imageSize)),
    transforms.ToTensor(),  # 先将PIL图像转为Tensor
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

步骤2：移除多余的二次转换逻辑 #

不需要自定义MyDataset和second_transform，直接基于Subset创建数据集后传入DataLoader即可：

train_set = datasets.ImageFolder(path, transform=train_transforms)

# 划分训练集与测试集
train_idx_, test_idx = train_test_split(
    np.arange(len(train_set.targets)),
    test_size=0.2,
    random_state=999,
    shuffle=True,
    stratify=train_set.targets
)
train_dataset_ = torch.utils.data.Subset(train_set, train_idx_)
test_dataset = torch.utils.data.Subset(train_set, test_idx)

# 从训练子集中划分验证集（注意：要用训练子集的标签做分层）
train_idx, val_idx = train_test_split(
    np.arange(len(train_dataset_)),
    test_size=0.15,
    random_state=999,
    shuffle=True,
    stratify=[train_set.targets[i] for i in train_idx_]
)
train_dataset = torch.utils.data.Subset(train_dataset_, train_idx)
valid_dataset = torch.utils.data.Subset(train_dataset_, val_idx)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=20) 
valid_loader = DataLoader(valid_dataset, batch_size=20) 
test_loader = DataLoader(test_dataset, batch_size=20)

# 验证加载是否正常
images, labels = next(iter(train_loader))

关键注意点 #

划分验证集时，不能直接使用train_dataset_.dataset.targets（这是原始完整数据集的标签），必须基于训练子集对应的标签列表做分层，才能保证划分后的数据类别分布与训练子集一致。

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