使用PyTorch将文件夹下的图片分为训练集和验证集

最新推荐文章于 2025-05-15 08:45:00 发布

原创

最新推荐文章于 2025-05-15 08:45:00 发布 · 1w 阅读

34 ·

CC 4.0 BY-SA版权

本文介绍如何使用PyTorch的ImageFolder类结合自定义代码，将文件夹中的图片数据集划分为训练集和验证集，便于模型训练过程中的数据验证。详细实现代码可在作者的Github仓库查阅。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

PyTorch提供了ImageFolder的类来加载文件结构如下的图片数据集：

root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png

root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png

使用这个类的问题在于无法将训练集（training dataset）和验证集（validation dataset）分开。我写了两个类来完成这个工作。

import os
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import ToTensor, Resize, Compose
from PIL import Image
from sklearn.model_selection import train_test_split

class ImageFolderSplitter:
    # images should be placed in folders like:
    # --root
    # ----root\dogs
    # ----root\dogs\image1.png
    # ----root\dogs\image2.png
    # ----root\cats
    # ----root\cats\image1.png
    # ----root\cats\image2.png    
    # path: the root of the image folder
    def __init__(self, path, train_size = 0.8):
        self.path = path
        self.train_size = train_size
        self.cl