Pytorch 利用imagefolder类划分train,test数据集

在torch.utils里有一个imagefolder类可以自动加载一个总文件夹下的各个类的图片和它的对应label,但是没有划分train和test的功能。另一个原因是在初始化的时候必须要声明transform,但是train和test的transform经常不同。

train_datasets = datasets.ImageFolder(train_dir)

train_size = int(0.8 * len(train_datasets))
test_size = len(train_datasets) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(train_datasets, [train_size, test_size])

train_dataset=train_dataset.dataset#这行很重要
for images, labels in train_dataset.imgs:
print(images)
print(labels)

经过random_split函数之后可以得到两个subset类,但是这个类和imagefolder初始化的对象不同。必须加上train_dataset=train_dataset.dataset才能保持和原来一样的数据格式一致。
这个时候再transform可以遍历train_dataset.imgs。不过images是路径,得用pil格式读取然后再transform。这样就可以不用再新建train,test文件夹而自动划分啦。

你可能感兴趣的:(Pytorch 利用imagefolder类划分train,test数据集)