图像二分类数据集划分和预处理以及批处理pytorch代码实现

最新推荐文章于 2025-10-24 15:08:00 发布

原创

最新推荐文章于 2025-10-24 15:08:00 发布 · 2.1k 阅读

28 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #人工智能 #大数据 #机器学习

本文介绍了如何在PyTorch中处理图像二分类问题，包括理解epoch、iteration和batch_size的关系，数据集的划分，自定义Dataset类，以及使用transforms进行图像增强和Dataloader进行批处理。示例中展示了未处理数据的存储结构，并提供了划分训练集、验证集和测试集的代码，以及实现Dataset类的关键方法。

在进行图像的数据处理之前我们先搞懂epoch、iteration和batch_size三者的关系

epoch:所有训练样本都送入到模型中称为一个eopch
iteration:一批样本输入到模型中，称之为一个iteration
batch_size:批大小，决定了一个epoch有多少个iteration
假设样本总数80 batch_size:8
则1epoch = 10iteration

下面先简单讲解一下思路
1.首先将图片进行训练集、验证集和测试集划分
2.接下来因为是使用pytorch的Dataloader类进行数据处理，因此我们要实现Dataset类，这个类需要自己实现
3.写完Dataset类可以用它来构建Dataloader类，用Dataloader类进行数据批处理,在进行数据批处理之前可以使用transforms对图像进行图像增强

首先先介绍未处理数据的文件存储结构.hello pytorch为项目的根路径
图像标签为1的存储路径：hello pytorch/data/RMB_data/1 #文件夹1下存放的是标签为1的图片数量100张
图像标签为100的存储路径：hello pytorch/data/RMB_data/100 # 文件夹100下存放的是标签为100的图片数量100张

在这里插入图片描述

划分训练集、验证集和测试集代码

import os
import random
import shutil

def makedir(new_dir): # 如果new_dir文件夹不存在则创造此文件夹
    if not os.path.exists(new_dir):
        os.makedirs(new_dir)

dataset_dir = os.path.join("..", "..", "data", "RMB_data")  # F:\pythonProject\hello pytorch\data\RMB_data
split_dir = os.path.join("..", "..", "data", "rmb_split") # 为划分的数据集准备的文件夹，此时还没有创建
train_dir = os.path.join(split_dir, "train") # 为训练集准备的文件夹，此时还没有创建
valid_dir = os.path.join(split_dir, "valid") # 为验证集创建的文件夹，此时还没有创建
test_dir = os.path.join(split_dir, "test") # 为测试集创造的文件夹，此时还没有创建

# 为测试集、训练集和验证集划分比例
train_pct = 0.8
valid_pct = 0.1
test_pct = 0.1

for root, dirs, files in os.walk(dataset_dir):

        for sub_dir in dirs:  # 1
            imgs = os.listdir(os.path.join(root, sub_dir)