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

最新推荐文章于 2024-04-29 22:01:40 发布

岁月风尘仆

最新推荐文章于 2024-04-29 22:01:40 发布

阅读量6k

点赞数 12

分类专栏：解决问题锦囊文章标签： pytorch

本文链接：https://blog.csdn.net/qq_36959832/article/details/118143485

版权

一、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-maste