pytorch学习1：如何加载自己的训练数据

转载自：https://blog.csdn.net/woshicao11/article/details/78318156，本文只做个人记录学习使用，版权归原作者所有。

Pytorch中文文档已出（http://pytorch-cn.readthedocs.io/zh/latest/）。第一篇博客献给了pytorch，主要是为了整理自己的思路。

原来使用caffe，总是要编译，经历了无数的坑。当开始接触pytorch时，果断拔草caffe。

学习Pytorch最好有一些深度学习理论基础才更好开，废话不多说，进入主题。

1 先有个框框，再往里面填东西
当训练一个神经网络的时候，我们需要有数据，有模型，并且需要设置训练的参数。为了不乱，我们最好分别定义三个文件，分别是：数据准备和预处理traindataset.py+编写模型model.py+如何训练main.py（xx.py，xx自己可任意取名）。

今天我们只讲数据准备与预处理阶段：traindataset.py（怎样命名无所谓，as you like）。这个文件的作用是什么呢？

统一将图像（或矩阵）返回成torch能处理的[original_iamges.tensor，label.tensor]
我们先跳跃一下看中文介绍是如何导入数据：

torch.utils.data.Dataloader(dataset, batch_size=1, shuffle=False, sample=None, num_workers=0, collate_fn=, pin_memory=False, drop_last=False)

我们一般关注dataloader的四个参数，dataset, batch_size, shuffle, num_workers=0

batch_size是你批处理数目，shuffle是否每个epoch都打乱，workers是载入数据的线程数。

我们具体再看看dataset——加载数据的数据集

这个dataset是[original_images.tensor, label.tensor]之类的，我们定义的traindataset.py就是产生这个dataset的，然后只需要在main.py文件import就可以调用！

from traindataset import *

2 定义一个自己的py文件产生自己的dataset

这个py文件一定要：

1）能输入自己的数据路径2）还能预处理

step 1:先导入你肯定需要的库路径

import torch.utils.data
import torch
from torchvision import transforms

torch.utils.data模块是子类化你的数据

transforms库对数据进行预处理

step 2:自定义dataset类（子类化你的数据）

class MyTrainData(torch.utils.data.Dataset)    

这里继承了torch.utils.data.Dataset这个类，我们看看这个类在中文文档中介绍：

所有其他类数据集都应该进行子类化。所有子类应该override_len_和_getitem_,前者提供了数据集的大小，后者支持整数索引，范围是从0到len(self),当然还有一个初始化_init_()_

类：属性+方法，_init_()就是定义自己的属性

我们脸谱化基础框架，再往里面加东西（以下为基础框架）

import torch.utils.data as data
import torch
import torchvision.transforms as transforms

class MyTrainData(torch.utils.data.Dataset) #子类化
    def __init__(self,root,transform=None,train=True): #第一步，初始化各个变量
        self.root=root
        self.train=train

    def __getitem__(self,idx): #第二步，装载数据，返回[img,label],idx就是一张纸读取
        
        img=imread(img_path) #img_path根据自己的数据自定义，灵活性很高
        img=torch.from_numpy(img).float() #需要转换成float
        
        gt=imread(gt_path) #读取gt，如果是分类问题，可以根据文件夹或者命名赋值0 1
        gt=torch.from_numpy(gt).float()

        return img,gt

    def __len__(self):
        return len(self.imagenumber)  #这个是必须返回的长度

现在往框框里填：

(1)是否transform如裁剪、归一化、旋转等？如果需要transform则要区分test和train，比如说我train需要随机旋转，但是test则不需要操作

(2)如何做到一张一张对应读取图片？可以自定义这些函数

以下贴出完整代码：

import torch.utils.data as data
import torch

from scipy.ndimage import imread
import os
import os.path
import glob

import torchvision.transforms as transforms

def make_dataset(root,train=True)：#读取自己的数据的函数
    dataset=[]
    if train:
        dirgt=os.path.join(root,'train_data/groundtruth')
        dirimg=os.path.join(root,'train_data/images')

    for fGT in glob.glob(os.path.join(dirgt,''*.jpg)):
        fName=os.path.basename(fGT)
        fImg='train_ori'+fName[8:]
        dataset.append([os.path.join(dirimg,fImg),os.path.join(dirgt,fName)]) 

return dataset

#自定义dataset的框架
class MyTrainData(data.Dataset):
    def __init__(self,root,transform=None,train=True): #初始化文件路径或文件名
        self.train=train
        if self.train:
            self.train_set_path=make_dataset(root,train)


    def __getitem__(self,idx):
        if self.train:
            img_path,gt_path=self.train_set_path[idx] 

            img=imread(img_path)
            img=np.atleast_3d(img).transpose(2,0,1).astype(np.float32)
            img=(img-img.min())/(img.max()-img.min())
            img=torch.from_numpy(img).float()

            gt=imread(gt_path)
            gt=np.atleast_3d(gt).transpose(2,0,1)
            gt=gt/255.0
            gt=torch.from_numpy(gt).float()

        return img,gt

    def __len__(self):
        return len(self.train_set_path)     

这里的py文件需要在最后的main.py中调用，所以root我并没有赋值，我会在main.py中赋值

这里我并没有用到“transform”进行预处理，如果你想用的话，在__getitem__()下面，return img,gt前面重新赋值

img=transforms.ToTensor(img)及gt=transforms.ToTensor(gt)

这里需要注意的是，查看中文文档transforms库有哪些变换，如果有需要涉及参数的如CenterCrop(size)，需要先实例化，如crop=transforms.CenterCrop(10)，再使用img=crop(img).