【pytorch基础】torchvision.ImageFolder()函数使用方法

转载自：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
)