【PyTorch】每个类对应一个文件夹的时候，PyTorch加载数据集

 PyTorch的ImageFolder是专门为分类提供的一个加载数据集的方法，使用方法如下：

"""
load the species file
"""

import torch
import torchvision
from torchvision import transforms, utils
import matplotlib.pyplot as plt

data_dir = 'D:\\c\\graduation\\train_data\\data_ex'

img_data = torchvision.datasets.ImageFolder(data_dir,
                                            transform=transforms.Compose([
                                                transforms.Resize(256),
                                                transforms.CenterCrop(224),
                                                transforms.ToTensor()])
                                            )

data_loader = torch.utils.data.DataLoader(img_data, batch_size=20, shuffle=True)

 

当需要划分数据集和验证集的时候，可以看【PyTorch】分类存储的图片，划分训练集与验证集（每个文件夹都存储同一类别的图片）这篇。

 

参考链接：

https://www.cnblogs.com/denny402/p/7512516.html