如何利用pytorch自定义一个数据集-创新互联

今天就跟大家聊聊有关如何利用pytorch自定义一个数据集，可能很多人都不太了解，为了让大家更加了解，小编给大家总结了以下内容，希望大家根据这篇文章可以有所收获。

自定义数据集

在训练深度学习模型之前，样本集的制作非常重要。在pytorch中，提供了一些接口和类，方便我们定义自己的数据集合，下面完整的试验自定义样本集的整个流程。

开发环境

• Ubuntu 18.04
• pytorch 1.0
• pycharm

实验目的

1. 掌握pytorch中数据集相关的API接口和类
2. 熟悉数据集制作的整个流程

实验过程

1.收集图像样本

以简单的猫狗二分类为例，可以在网上下载一些猫狗图片。创建以下目录：

• data-------------根目录
• data/test-------测试集
• data/train------训练集
• data/val--------验证集

在test/train/val之下在校分别创建2个文件夹，dog, cat

cat, dog文件夹下分别存放2类图像：

标签

之后写一个简单的python脚本，生成txt文件，用于指明每个图像和标签的对应关系。

格式: /cat/1.jpg 0 \n dog/1.jpg 1 \n .....

如图：

至此，样本集的收集以及简单归类完成，下面将开始采用pytorch的数据集相关API和类。

2. 使用pytorch相关类，API对数据集进行封装

2.1 pytorch中数据集相关的类，接口

pytorch中数据集相关的类位于torch.utils.data package中。

https://pytorch.org/docs/stable/data.html

本次实验，主要使用以下类：

torch.utils.data.Dataset
torch.utils.data.DataLoader

Dataset类的使用： 所有的类都应该是此类的子类(也就是说应该继承该类)。 所有的子类都要重写(override) __len()__, __getitem()__ 这两个方法。

这里和Java抽象类很相似，在抽象类abstract class中，一般会定义一些抽象方法abstract method,抽象方法：只有方法名没有方法的具体实现。如果一个子类继承于该抽象类，要重写(overrode)父类的抽象方法。

DataLoader类的使用：

2.2 实现

使用到的python package

源码

导入python包

import numpy as np
from skimage import io
from skimage import transform
import matplotlib.pyplot as plt
import os
import torch
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms
from torchvision.utils import make_grid