转载自:https://blog.csdn.net/c2250645962/article/details/105291782
背景
在使用pytorch做深度学习任务的数据加载时,常用的方式是使用torchvision.Dataset类定义数据读取,然后使用torch.utils.data.DataLoader定义数据加载器。该部分内容见 链接
不过,有些分类数据的文件目录组织形式如下:
即默认你的数据集已经自觉按照要分配的类型分成了不同的文件夹,一种类型的文件夹下面只存放一种类型的图片。
这时候,定义数据读取时,使用 torchvision包中的ImageFolder类会比Dataset类会更方便。
具体阐述
CLASS torchvision.datasets.ImageFolder(root, transform=None, target_transform=None, loader=<function default_loader>, is_valid_file=None)
一个通用数据加载器,其中图像以这种方式排列:
root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png
root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png
测试如下:
import torchvision.datasets as dset
dataset = dset.ImageFolder('./data/dogcat_2') #没有transform,先看看取得的原始图像数据
print(dataset.classes) #根据分的文件夹的名字来确定的类别
print(dataset.class_to_idx) #按顺序为这些类别定义索引为0,1...
print(dataset.imgs) #返回从所有文件夹中得到的图片的路径以及其类别
输出为
['cat', 'dog']
{'cat': 0, 'dog': 1}
[('./data/dogcat_2/cat/cat.12484.jpg', 0), ('./data/dogcat_2/cat/cat.12485.jpg', 0), ('./data/dogcat_2/cat/cat.12486.jpg', 0), ('./data/dogcat_2/cat/cat.12487.jpg', 0), ('./data/dogcat_2/dog/dog.12496.jpg', 1), ('./data/dogcat_2/dog/dog.12497.jpg', 1), ('./data/dogcat_2/dog/dog.12498.jpg', 1), ('./data/dogcat_2/dog/dog.12499.jpg', 1)]
具体调用代码
# 指定读取的图片路径
train_root = './train/
# transform函数组合
train_transform = transforms.Compose([
transforms.Resize(224),
transforms.RandomResizedCrop(224,scale=(0.6,1.0),ratio=(0.8,1.0)),
transforms.RandomHorizontalFlip(),
torchvision.transforms.ColorJitter(brightness=0.5, contrast=0, saturation=0, hue=0),
torchvision.transforms.ColorJitter(brightness=0, contrast=0.5, saturation=0, hue=0),
transforms.ToTensor(),
transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
# 使用ImageFolder读取数据
all_data = torchvision.datasets.ImageFolder(
root=train_root,
transform=train_transform
)
# 定义数据加载器
train_set = torch.utils.data.DataLoader(
all_data,
batch_size=BTACH_SIZE,
shuffle=True
)