PyTorch学习笔记（二）划分数据集

最新推荐文章于 2024-08-27 22:36:12 发布

longrootchen

最新推荐文章于 2024-08-27 22:36:12 发布

阅读量4.8k

点赞数 1

分类专栏： PyTorch学习笔记文章标签：人工智能计算机视觉 pytorch 深度学习

本文链接：https://blog.csdn.net/longrootchen/article/details/105447058

版权

本文介绍了在PyTorch中如何处理计算机视觉任务的数据集，包括将图像按标签分类放入不同子目录以及如何划分训练、验证和测试集。通过编写categorise.py和split.py脚本，实现了数据的正确组织和划分，确保符合PyTorch的读取需求。

摘要由CSDN通过智能技术生成

Environment

OS: macOS Mojave
Python version: 3.7
PyTorch version: 1.4.0
IDE: PyCharm

0. 写在前面

对于计算机视觉深度学习任务，对数据的处理包括

划分数据集，将数据集划分为划分训练集、验证集、测试集；
预处理，对图像进行数据增强和标准化；
读取，读取一个 batch 的数据输入模型。

PyTorch 读取数据进行训练时，要求数据按照特定的目录结构放好，所以划分数据集就是要将数据整理为特定的目录格式。

1. 将图像按标签分类放入不同子目录

到手的训练图像数据集有时候并未是方便 PyTorch 读取的目录格式。以 TinyMind人民币面值识别任务的训练集为例，共有 39620 张图片，train_face_value_label.csv 中为每张图片对应的标签信息

├── train
│   ├── 39620 images.jpeg
└── train_face_value_label.csv

编写一个 categorise.py 按类别存放到不同目录中，形成以下目录结构

├── categorise.py
├── train
│   ├── 0.1
│	│	└── 4233 images.jpg
│	├── 0.2
│	│	└── 4373 images.jpg
│	├── 0.5
│	│	└── 4407 images.jpg
│	├── 1.0
│	│	└── 4424 images.jpg
│	├── 2.0
│	│	└── 4411 images.jpg
│	├── 5.0
│	│	└── 4413 images.jpg
│	├── 10.0
│	│	└── 4283 images.jpg
│	├── 50.0
│	│	└── 4408 images.jpg
│	├── 100.0
│	│	└── 4668 images.jpg
└── train_face_value_label.csv

import os
import time
import shutil
import pandas as pd

label_path = os.path.join(os.curdir, 'train_face_value_label.csv')
labels = pd.read_csv(label_path)

# move each image to the specified-class dir
since = time.time()
data_dir = os.path.join(os.curdir