pytorch 加载自制图像数据集并划分训练集和测试集：

一、pytorch 加载自制图像数据集并划分训练集和测试集步骤： 

1、对数据集进行文件夹组织处理以适应ImageFolder的要求，其格式为： 
        datasrc/dog/xxx.png 
        datasrc/dog/xxy.png 
        datasrc/dog/xzz.png 

        datasrc/cat/xxx.png 
        datasrc/cat/xxy.png 
        datasrc/cat/xzz.png 

2、然后将ImageFolder的返回结果作为DataLoader的输入，完成训练集数据和测试集数据的迭代。         torchvision.datasets.ImageFolder()读取图像
       

3、 random_split实现数据集的划分， 
        train, test = torch.utils.data.random_split(full_dataset, [train_lengths, test_lengths]) 
        train_lengths, test_lengths是一个list，按照对应的数量返回数据个数。数据量总和等于all_dataset中的数据个数，不是按比例划分的。 

4 使用DataLoader完成训练集数据和测试集数据的读取迭代。 
  torch.utils.data.DataLoader()加载数据集 

二、代码实例

# 工具类
import torch
import torchvision
from torch.utils.data import DataLoader, random_split
from torch.utils.tensorboard import SummaryWriter


data_path = 'E:/AI_Project/pytorch_B_zhan/garbage-classification-master/dataset-resized/'

full_data = torchvision.datasets.ImageFolder(root=data_path,transform=torchvision.transforms.ToTensor()) #读取图像，将numpy的ndarray或PIL.Image读的图片转换成形状为(C,H, W)的Tensor格式，且/255归一化到[0,1.0]之间
train_size = int(len(full_data)*0.8)  #这里train_size是一个长度矢量，并非是比例，我们将训练和测试进行8/2划分
test_size = len(full_data) - train_size


train_dataset, test_dataset =random_split(full_data, [train_size, test_size])

train_loader = DataLoader(train_dataset, batch_size=1, num_workers=0, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1, num_workers=0, shuffle=False)
'''
SummaryWriter查看数据集的划分
1、writer = SummaryWriter("logs")将log存在文件夹logs下。
 
'''
writer = SummaryWriter("logs")


for epoch in range(3):# 查看前3个epoch图像
    step = 0
    for data in train_loader:
        imgs, targets = data #读取图像和标签
        print(imgs.shape)
        print(targets)
        writer.add_images("Epoch: {}".format(epoch), imgs, step)
        writer.add_scalar("Epoch: {}".format(epoch), targets, step)
        step = step + 1

writer.close()

三、效果图