将自己的数据集引入Pytorch是搭建属于自己的神经网络的重要一步，这里我设计了一个简单的实验，结合这个实验代码，我将逐句教会大家如何将数据引入DataLoader。

这里以目标检测为例，一个batch中包含图片文件、先验框的框体坐标、目标类型，相对而言更加全面。大家亦可根据自己的数据结构和需求进行修改。

一、数据文件分析

标准的Voc格式是无法直接注入模型的，而如果在训练程序中进行处理即拖慢了运算速度，又难以保证数据集分割的一致性。最好是使用一个独立程序完成数据集的分割、组织、暂存。这里参考了Bubbliiiing的做法，将数据集信息暂存为txt文件。其中一条具体数据的格式如下：

../VOC2007/JPEGImages/0.jpg 166,121,336,323,0 1052,372,1371,924,1
#   文件路径（绝对路径为佳）/ # 先验框框体信息1/  # 先验框框体信息1/

文件路径和框体信息之间采用空格分开；框体信息内部以逗号分开，前4个为坐标信息，最后一个为分类信息。

在随后的程序中，我们将循环读取这个文件中的数据来获取数据集信息。

二、载入数据

1.定义DataSet超参

在开始重写DataSet前，我们需要定义一些用来控制DataSet的参数。

#----自定义DataSet，继承自torch.utils.data.DataSet
class MyDataSet(Dataset):#----参数定义，输入的参数分别为数据行、输入图像尺寸，类型数def __init__(self,file_Lines,inp_Shape,num_Classes):super(MyDataSet,self).__init__()#  将局部形参变为类内的全局变量self.length = len(file_Lines)   #将文件数赋值给lengthself.file_Lines = file_Linesself.inp_Shape = inp_Shapeself.num_Classes = num_Classes

2.重写len函数

没什么技巧，因为我们刚刚将数据行(file_Lines)的长度赋值给了self.length，直接返回这个值就能拿到数据集的长度了。这也是设置超参数的意义所在。

def __len__(self):return self.length

3.重写getitem函数

此函数每次会获取单个文件，DataLoader通过反复调用这个函数最终获取整个数据集，我们写入的，这个函数自带一个index用于控制获取的行数。

①分解数据集文件

按照我们上面解析的文件格式，我们用split函数分割index行的文本（空格分割），得到的file_item中第1个元素为文件的绝对路径，后续元素为目标先验框的信息。

    def __getitem__(self, index):index = index % self.length     #计算batch长度file_item = self.file_Lines[index].split()  #按空格拆分文件行，其中的元素分别为：文件路径、先验框坐标(n个)

file_item的值如下图所示：

②加载图片文件

同样没什么技巧，拿到文件路径后直接打开就好了。需要注意的是神经网络的输入需要为固定的形状（图片尺寸和颜色通道数），如果图片为灰阶图则需要将其颜色通道扩充为3个（RGB图）

        img = Image.open(file_item[0])              #打开图片img = cvt2RGB(img)                          #若图像为灰度图需要先转换为RGB图（神经网络输入为3通道）

#----将图像转换为RGB----#
def cvt2RGB(img):if len(np.shape(img)) == 3 and np.shape(img)[2]==3:return img                                      #为RGB不需要转换else:img = img.convert('RGB')return img

③拆分框体信息

同样同样没什么好说的，遍历分割file_item从1开始的元素就好了

box_info = np.array([np.array(list(map(int,box.split(',')))) for box in file_item[1:]])   
#从文件中加载先验框坐标和类型(从第1个元素开始)

④将图片变形

这一步也不是必须的，可以选择在开始训练之前对图片信息进行处理。但是在程序中处理需要注意一点，在改变图像的同时需要以同样的比例改变先验框的坐标。

img,new_box = self.resize_img_withBox(img,box_info,self.inp_Shape)

这里给出一个无损变换大小的函数，若不指定参数则不变化。

    def resize_img_withBox(self,img, box, size=[0,0]):      #输入参数分别为：原图、先验框列表、变形后的图片大小iw,ih = img.sizew,h = sizenew_box=[]#  若没有指定大小则不需要变形，若指定了大小则进行变形if size!=(0,0):scale = min(w/iw,h/ih)                          #获取变形比例nw = int(iw*scale)                              #计算变形后的长宽nh = int(ih*scale)dx = (w-nw)//2dy = (h-nh)//2#  图像变形img = img.resize((nw,nh),Image.BICUBIC)new_img = Image.new('RGB',size,(128,128,128))   #创建一张灰色背景new_img.paste(img,((w-nw)//2,(h-nh)//2))        #将变形后的图片贴进背景中央#  先验框变形if len(box)>0:np.random.shuffle(box)box[:, [0,2]] = box[:, [0,2]]*nw/iw + dxbox[:, [1,3]] = box[:, [1,3]]*nh/ih + dybox[:, 0:2][box[:, 0:2]<0] = 0box[:, 2][box[:, 2]>w] = wbox[:, 3][box[:, 3]>h] = hbox_w = box[:, 2] - box[:, 0]box_h = box[:, 3] - box[:, 1]new_box = box[np.logical_and(box_w>1, box_h>1)] # discard invalid boxreturn new_img,new_box

⑤变换图片的通道

标准的图片通道为RGB，而在pytorch中图片的通道为BGR，所以我们需要对通道进行调整，同时为其附加batch通道。

img = np.transpose(preprocess_input(np.array(img, dtype=np.float32)), (2, 0, 1))

前面的函数是一个增强函数，会给RGB三个通道加上不同的权值，至于权值则是一个默认权值（我也不知道为什么用这个）

#----为图像加权，这是一般默认的参数----#
def preprocess_input(image):image   = np.array(image,dtype = np.float32)[:, :, ::-1]mean    = [0.40789655, 0.44719303, 0.47026116]std     = [0.2886383, 0.27408165, 0.27809834]return (image / 255. - mean) / std

⑥拆分坐标信息和分类信息

如题，将坐标信息和分类信息从先验框信息组中进行分割

        #  拆分先验框坐标和类型box_data = np.zeros((len(new_box),5))if(len(box_info)>0):box_data[:len(box_info)] = box_infobox = box_data[:,:4]label = box_data[:,-1]

完成上述步骤后，将得到的数据返回即可，完整的getitem函数如下：

    def __getitem__(self, index):index = index % self.length     #计算batch长度#  读取文件file_item = self.file_Lines[index].split()  #按空格拆分文件行，其中的元素分别为：文件路径、先验框坐标(n个)img = Image.open(file_item[0])              #打开图片img = cvt2RGB(img)                          #若图像为灰度图需要先转换为RGB图（神经网络输入为3通道）box_info = np.array([np.array(list(map(int,box.split(',')))) for box in file_item[1:]])   #从文件中加载先验框坐标和类型(从第1个元素开始)#  对图像进行变形(含先验框变形)img,new_box = self.resize_img_withBox(img,box_info,self.inp_Shape)#  将图像进行加权#img = np.transpose(preprocess_input(np.array(img, dtype=np.float32)), (2, 0, 1))img = np.transpose(np.array(img))#  拆分先验框坐标和类型box_data = np.zeros((len(new_box),5))if(len(box_info)>0):box_data[:len(box_info)] = box_infobox = box_data[:,:4]label = box_data[:,-1]return img,box,label

三、数据封包

我们在训练时肯定不能这样一个一个训练，一般情况我们训练时会将这些数据打包成一个个的patch送给迭代器，而collate_fn就是做这个的，需要注意collate_fn并不是DataSet类的成员。

这个函数使dataloader自动使用的，其中的images、bboxes、labels 将会在训练过程中用到，这里我们只要确保将数据装入对应的容器中即可。

# DataLoader中collate_fn使用
def my_collate(batch):images = []bboxes = []labels = []for img, box, label in batch:images.append(img)bboxes.append(box)labels.append(label)images = np.array(images)return images, bboxes, labels

四、调用

①读取数据集文件(txt)

    file_path = "数据集文件的路径"with open(file_path) as f:train_lines = f.readlines()

②实例化DataSet

    train_dataset = MyDataSet(train_lines,input_shape,num_classes)train_loader = DataLoader(train_dataset, shuffle = True, batch_size = 32, num_workers = 1, collate_fn=my_collate)

其中num_workers是线程数；batch_size是单个batch的大小；collate_fn指向我们刚刚重写的collate_fn；shuffle表示是否打乱数据集的顺序。

③训练

当然，这里不是真的训练，我们用一个展示函数来代替训练。

    print("开始打印结果")for item in train_dataset:img, box, label = itemimg = img.transpose(1,2,0)print(img.shape)im = Image.fromarray(np.uint8(img))im.show()input("按任意键继续")