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

1. 引言

在使用深度学习进行图像分类任务时，常常需要将数据集划分为训练集和验证集。训练集用于模型的训练，验证集用于评估模型的性能。

2. 准备工作

在开始编写代码前，我们需要先准备好数据集的文件夹结构。假设我们的数据集文件夹为"data"，其中包含多个类别的子文件夹，每个子文件夹中包含对应类别的图片。

data/

    ├── class1/
    │   ├── image1.jpg
    │   ├── image2.jpg
    │   ├── ...
    ├── class2/
    │   ├── image1.jpg
    │   ├── image2.jpg
    │   ├── ...
    ├── class3/
    │   ├── image1.jpg
    │   ├── image2.jpg
    │   ├── ...
    └── ...

3. 代码实现

3.1 导入模块

import os

import random
import shutil

3.2 定义函数

我们首先定义一个函数来划分数据集，并将训练集和验证集分别存放在不同的文件夹中。

def split_dataset(data_dir, train_ratio):

    # 创建训练集和验证集的文件夹
    os.mkdir(os.path.join(data_dir, 'train'))
    os.mkdir(os.path.join(data_dir, 'val'))
    
    # 获取数据集中的类别名
    class_names = os.listdir(data_dir)
    
    for class_name in class_names:
        # 获取当前类别的图片路径列表
        image_paths = [os.path.join(data_dir, class_name, img) for img in os.listdir(os.path.join(data_dir, class_name))]
        
        # 打乱图片路径列表的顺序
        random.shuffle(image_paths)
        
        # 计算训练集的数量
        train_size = int(len(image_paths) * train_ratio)
        
        # 将图片拷贝到训练集文件夹
        for img_path in image_paths[:train_size]:
            shutil.copy(img_path, os.path.join(data_dir, 'train', class_name))
        
        # 将图片拷贝到验证集文件夹
        for img_path in image_paths[train_size:]:
            shutil.copy(img_path, os.path.join(data_dir, 'val', class_name))

上述代码中，data_dir表示数据集的文件夹路径，train_ratio表示训练集所占比例。例如，如果train_ratio=0.8，则训练集占整个数据集的80%。

3.3 调用函数

接下来，我们调用刚刚定义的函数来实现数据集的划分。

data_dir = 'data'

train_ratio = 0.8
split_dataset(data_dir, train_ratio)

运行上述代码后，数据集将被划分为训练集和验证集，分别存放在"data/train"和"data/val"文件夹中。

4. 结语

在本文中，我们介绍了如何使用PyTorch将文件夹下的图片分为训练集和验证集。首先，我们准备好数据集的文件夹结构，然后编写代码实现数据集的划分。最后，我们调用刚刚定义的函数来实现数据集的划分。通过合理划分数据集，我们能够更好地评估模型的性能，并进行更有效的模型训练。