用于提供多类图像数据集的方法，其中文件夹名称可以用作Pytorch中的标签？

Posted 2023-01-31

我想在Pytorch中提供多类图像数据集，在数据集的主文件夹中，我还有15个具有不同名称的文件夹，我想使用文件夹名称作为标签。例如，一个文件夹名称是Airplanes并包含图像（1245张图像），另一个文件夹名称是Cars并包含Cars图像（997），同样，每个文件夹具有不同数量的图像。现在，我想加载它们来训练我的模型并对其进行测试，但是我没有用于训练和测试的单独文件夹。我想使用文件夹名称作为标签，还希望将数据集按相等的比例分为训练和测试。在这种情况下，您的指导将不胜感激。谢谢

答案

要将数据集分为训练和测试数据集，可以使用random_split函数：

import torch
from torchvision import datasets, transforms
from torch.utils import data
import numpy as np

dataset = datasets.ImageFolder('path_to_dataset', transform=transforms.ToTensor())

lengths = [int(np.ceil(0.5*len(dataset))),
           int(np.floor(0.5*len(dataset)))]
train_set, test_set = data.random_split(dataset, lengths)

train_dataloader = data.DataLoader(train_set, batch_size=...)
test_dataloader = data.DataLoader(test_set, batch_size=...)

如果要对训练和测试数据集执行单独的转换，请看这里：How to use different data augmentation for Subsets in PyTorch