torchvision.datasets.ImageFolder用法

torchvision之ImageFolder用法

当前文件夹：

/root
-- train/
---- 001.jpg
---- 002.jpg
---- ....
-- test/
---- 001.jpg
---- 002.jpg
---- ....

导入当前文件夹的训练集，用下面的代码会报错

root = "./root/train"
train_dataset = ImageFolder(root = root_path)

正确的路径

data/
├── train/
|   ├── class_0/
|   |   ├── 001.jpg
|   |   ├── 002.jpg
|   |   └── 003.jpg
|   └── class_1/
|       ├── 004.jpg
|       └── 005.jpg
└── test/
    ├── class_0/
    |   ├── 006.jpg
    |   └── 007.jpg
    └── class_1/
        ├── 008.jpg
        └── 009.jpg

再用下面的代码，能正确导入训练集

train_dataset = ImageFolder(root='data/train')
test_dataset  = ImageFolder(root='data/test')

得到的数据类型可以放入dataloader中

train_dataset = ImageFolder(root='data/train')
torchvision.datasets.folder.ImageFolder
train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size, shuffle=True,
                                               num_workers=nw)

官方给出示例

        root/dog/xxx.png
        root/dog/xxy.png
        root/dog/[...]/xxz.png

        root/cat/123.png
        root/cat/nsdf3.png
        root/cat/[...]/asd932_.png

** 注意两个子文件夹中的[…]**
此外：train_dataset.class_to_idx能返回一级文件夹的索引
可以使用len（train_dataset）查看图片数量

报错
Found 0 files in subfolders of: C:\Users\xxx\Desktop\test
Supported extensions are: .jpg,.jpeg,.png,.ppm,.bmp,.pgm,.tif,.tiff,.webp