Pytorch加载自己的图片分类数据集（torchvision.datasets.ImageFolder使用详解（数据加载器））

假设数据集文件夹与你的.py文件同路径，如下结构,在这里，我是一个二分类任务的数据集：

• ├─ project
•         │ ├─ main.py
•         │ ├─ dataSet
•                 │ │ ├─ train
•                          ├─ ants
•                          ├─ bees
•                 │ │ ├─ val
•                          ├─ ants
•                          ├─ bees

from torchvision.datasets import ImageFolder
from torchvision import transforms

#加上transforms
normalize=transforms.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])
transform=transforms.Compose([
    transforms.RandomCrop(180),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(), #将图片转换为Tensor,归一化至[0,1]
    normalize
])

train_dataset=ImageFolder('./data/train/',transform=transform)

print(train_dataset.classes)  #根据分的文件夹的名字来确定的类别
print(train_dataset.class_to_idx) #按顺序为这些类别定义索引为0,1...
print(train_dataset.imgs) #返回从所有文件夹中得到的图片的路径以及其类别
print(train_dataset[0])#dataset,它的结构就是[(img_data,class_id),(img_data,class_id),…],下面我们打印第一个元素：

['ants', 'bees']

{'ants': 0, 'bees': 1}

[('./data/train/ants/0013035.jpg', 0),

('./data/train/ants/1030023514_aad5c608f9.jpg', 0),

...

('./data/train/bees/1092977343_cb42b38d62.jpg', 1), ('./data/train/bees/1093831624_fb5fbe2308.jpg', 1), ]

(tensor([[[-0.3255, -0.3333, -0.3490,  ..., -0.3961, -0.3882, -0.3882],

         [-0.3333, -0.3333, -0.3412,  ..., -0.3961, -0.3961, -0.3961],

         [-0.3412, -0.3412, -0.3333,  ..., -0.4039, -0.4039, -0.4039],

         ...,

         [-0.3176, -0.3098, -0.3098,  ..., -0.2863, -0.2627, -0.2627],

         [-0.3176, -0.3020, -0.3020,  ..., -0.3020, -0.2863, -0.2706],

         [-0.3176, -0.3020, -0.3020,  ..., -0.3255, -0.3412, -0.3020]],

        [[ 0.2078,  0.2000,  0.1843,  ...,  0.1608,  0.1686,  0.1686],

         [ 0.2000,  0.2000,  0.1922,  ...,  0.1608,  0.1608,  0.1608],

         [ 0.1922,  0.1922,  0.2000,  ...,  0.1529,  0.1529,  0.1529],

         ...,

         [ 0.2549,  0.2627,  0.2627,  ...,  0.2392,  0.2235,  0.1608],

         [ 0.2549,  0.2706,  0.2706,  ...,  0.2863,  0.2863,  0.2549],

         [ 0.2549,  0.2706,  0.2706,  ...,  0.3020,  0.2863,  0.2627]],

        [[ 0.8431,  0.8353,  0.8196,  ...,  0.8039,  0.8118,  0.8118],

         [ 0.8353,  0.8353,  0.8275,  ...,  0.8039,  0.8039,  0.8039],

         [ 0.8275,  0.8275,  0.8353,  ...,  0.7961,  0.7961,  0.7961],

         ...,

         [ 0.8588,  0.8667,  0.8667,  ...,  0.8510,  0.8118,  0.7255],

         [ 0.8588,  0.8745,  0.8745,  ...,  0.8196,  0.8275,  0.8118],

         [ 0.8588,  0.8745,  0.8745,  ...,  0.7804,  0.7804,  0.8275]]]), 0)

参考资料

//可以说真的讲的很好

pytorch学习笔记七：torchvision.datasets.ImageFolder使用详解_耐心的小黑的博客-CSDN博客_datasets.imagefolder

pytorch之ImageFolder使用详解

https://www.jb51.net/article/177931.htm

pytorch实践（更改resnet50全连接层，重新训练，用于二分类）_Miaodawang的博客-CSDN博客_resnet全连接层