读取并预处理自己数据集的一种方式

引入必要的包

from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt
from torchvision import transforms as tfs
from torch.utils.data import DataLoader
import numpy as np
import torch

加载我们的图片数据集，路径为整个训练集文件夹的路径

folder_set = ImageFolder('./data/hymenoptera_data/train')

查看名称和类别下标的对应

print(folder_set.class_to_idx)

得到所有图片的名字和标签
print(folder_set.imgs)

取出其中一个数据并显示图片

im, label = folder_set[0]
plt.imshow(im)
plt.show()
print(label)

传入数据预处理方式
定义预处理函数：这里处理了图片大小，并将其转换成了Tensor

def data_tf(x):
    img_aug = tfs.Compose([
        tfs.Resize((32, 32))
    ])
    x = img_aug(x)
    x = torch.Tensor(np.array(x))
    return x

加载数据集并查看

folder_set = ImageFolder('./data/hymenoptera_data/train',transform=data_tf)
im, label = folder_set[0]
print(im.shape)

将五个数据作为一个batch

train_data1 = DataLoader(folder_set, batch_size=5, shuffle=True)
for im, label in train_data1:
    print(im)