用torchvision.datasets.ImageFolder加载图片数据集

一、项目结构

二、代码

data_loader = torch.utils.data.DataLoader(
    torchvision.datasets.ImageFolder('traing_dataset',
       transform=torchvision.transforms.Compose([
                torchvision.transforms.Resize([28, 28]),       # 裁剪图片
                torchvision.transforms.Grayscale(1),           # 单通道
                torchvision.transforms.ToTensor(),             # 将图片数据转成tensor格式
                torchvision.transforms.Normalize(              # 归一化
                     (0.1307,), (0.3081,))
                ])),
    batch_size=10, shuffle=False)　　　　　　　　　　　　　　　　　　# 10张图片

三、显示效果

def plot_image(img, label, name):
    fig = plt.figure()
    for i in range(6):　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　# 只显示6张
        plt.subplot(2, 3, i+1)　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　# 2行3列第i+1张
        plt.tight_layout()
        plt.imshow(img[i][0]*0.3081+0.1307, cmap='Greys', interpolation='none')
        plt.title("{}:{}".format(name, label[i].item()))　　　　　　　　　　　　　　　　# 标题名称
        plt.xticks([])
        plt.yticks([])
    plt.show()

x, y = next(iter(data_loader))　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　　# 文件夹的名称即为图片的label
print(x.shape, y.shape)
plot_image(x, y, 'image')