[Pytorch]DataSet和DataLoader逐句详解

        将自己的数据集引入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)   #将文件数赋值给length
        self.file_Lines = file_Lines
        self.inp_Shape = inp_Shape
        self.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.size
        w,h = size
        new_box=[]
        #  若没有指定大小则不需要变形,若指定了大小则进行变形
        if size!=(0,0):
            scale = min(w/iw,h/ih)                          #获取变形比例
            nw = int(iw*scale)                              #计算变形后的长宽
            nh = int(ih*scale)
            
            dx = (w-nw)//2
            dy = (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 + dx
                box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
                box[:, 0:2][box[:, 0:2]<0] = 0
                box[:, 2][box[:, 2]>w] = w
                box[:, 3][box[:, 3]>h] = h
                box_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 box
        return 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_info
        box = 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_info
        box = 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 = item
        img = img.transpose(1,2,0)
        print(img.shape)
        im = Image.fromarray(np.uint8(img))
        im.show()
        input("按任意键继续")

         这个昆虫数据集太恶心了就不给大家看了

猜你喜欢

转载自blog.csdn.net/weixin_37878740/article/details/129350390