欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

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

程序员文章站 2022-03-10 13:27:48
...
一、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/data