pytorch dataloader详解

构建自己的dataloader是模型训练的第一步,本篇文章介绍下pytorch与dataloader以及与其相关的类的用法。

DataLoader类中有一个必填参数为dataset,因此在构建自己的dataloader前,先要定义好自己的Dataset类。这里先大致介绍下这两个类的作用:

  • Dataset:真正的“数据集”,它的作用是:只要告诉它数据在哪里(初始化),就可以像使用iterator一样去拿到数据,继承该类后,需要重载__len__()以及__getitem__
  • DataLoader:数据加载器,设置一些参数后,可以按照一定规则加载数据,比如设置batch_size后,每次加载一个batch_siza的数据。它像一个生成器一样工作。

有小伙伴可能会疑惑,自己写一个加载数据的工具似乎也没有多“困难”,为何大费周章要继承pytorch中类,按照它的规则加载数据呢?关于这点可以参考这里:pytorch dataloader,总结一下就是:

  • 当数据量很大的时候,单进程加载数据很慢
  • 一次全加载过来,会占用很大的内存空间(因此dataloader是一个生成器,惰性加载)
  • 在进行训练前,往往需要一些数据预处理或数据增强等操作,pytorch的dataloader已经封装好了,避免了重复造轮子

一、使用方法

两步走:

  1. 定义自己的Dataset类,具体要做的事:

    • 告诉它去哪儿读数据,并将数据resize为统一的shape(可以思考下为什么呢)
    • 重写__len__()以及__getitem__,其中__getitem__中要确定自己想要哪些数据,然后将其return出来。
  2. 将自己的Dataset实例传到Dataloder中并设置想要的参数,构建自己的dataloader

下面简单加载一个目录下的图片以及label:

import os
import numpy as np

from torch.utils.data.dataset import Dataset
from torch.utils.data.dataloader import DataLoader
import cv2

# Your Data Path
img_dir = '/home/jyz/Downloads/classify_example/val/骏马/'
anno_file = '/home/jyz/Downloads/classify_example/val/label.txt'


class MyDataset(Dataset):
    def __init__(self, img_dir, anno_file, imgsz=(640, 640)):
        self.img_dir = img_dir
        self.anno_file = anno_file
        self.imgsz = imgsz
        self.img_namelst = os.listdir(self.img_dir)

    # need to overload
    def __len__(self):
        return len(self.img_namelst)

    # need to overload
    def __getitem__(self, idx):
        with open(self.anno_file, 'r') as f:
            label = f.readline().strip()
        img = cv2.imread(os.path.join(img_dir, self.img_namelst[idx]))
        img = cv2.resize(img, self.imgsz)
        return img, label


dataset = MyDataset(img_dir, anno_file)
dataloader = DataLoader(dataset=dataset, batch_size=2)

# display
for img_batch, label_batch in dataloader:
    img_batch = img_batch.numpy()
    print(img_batch.shape)
    # img = np.concatenate(img_batch, axis=0)
    if img_batch.shape[0] == 2:
        img = np.hstack((img_batch[0], img_batch[1]))
    else:
        img = np.squeeze(img_batch, axis=0)  # 最后一张图时,删除第一个维度
    print(img.shape)
    cv2.imshow(label_batch[0], img)
    cv2.waitKey(0)

上面是一次加载两张图片,效果如下:
在这里插入图片描述其实从这里可以看出,为什么要在Dataset中将数据resize为统一的shape。因为dataloader加载数据时,将一个batch_size的数据拼接成一个大的tensor,如果shape不同,就无法拼接了。就像这两张图片加入shape不一样就无法通过拼接的方式show出来一样。

二、结论

  1. 使用pytorch的dataloader,需要先构建自己的Dataset
  2. 构建自己的Dataset,需要重载__len__()以及__getitem__
  3. 数据地址:example data,提取码: a1ds

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
扎眼的阳光的头像扎眼的阳光普通用户
上一篇 2023年3月5日 下午3:30
下一篇 2023年3月5日

相关推荐