pytorch之数据集

在开始深度学习之前，除了需要构建合适的网络结构外，还需要有一个合适的数据集，一个出色的数据集+一个出色的神经网络模型才能有令人满意的结果。倘若我们已经找到了一份比较合适的数据集，就需要对数据集操作提取相应的训练集、验证集和测试集信息。
在pytorch中提供了两个与提取数据集信息密切相关的类，一个是torch.utils.data.Dataset，另一个是torch.utils.data.DataLoader。
torch.utils.data.DataLoader可以多线程的加载数据集，具体使用为：torch.utils.data.DataLoader(dataset, batch_size=args.batchSize, shuffle=True, num_works=args.nThreads，pin_memory=False)，其中dataset为需要被加载的数据，batch_size是批量操作的大小，shuffle表示要不要打乱数据，num_works是几个线程，pin_memory表示要不要锁页内存，一般在GPU上训练时设置成True，在CPU上训练时设置成False。
torch.util.data.Dataset是一个抽象类，每个它的子类都应该重写它未实现的方法：

class Dataset(object):
    """An abstract class representing a Dataset.

    def __getitem__(self, index):
        raise NotImplementedError

    def __len__(self):
        raise NotImplementedError

    def __add__(self, other):
        return ConcatDataset([self, other])

具体说来，对于Dataset的一个子类，至少需要实现init()、getitem()和len()方法，其中init()用来初始化子类，getitem()用来返回每一个数据项，len用来返回该数据集的大小。
下面是一个具体的在自己的数据集上操作的例子（代码不是自己写的，我是学习的时候写这篇博客来记录的，如果本人看到觉得侵权的话，可以联系我，我会第一时间处理~）：

class Garbage_Loader(Dataset):
    def __init__(self, txt_path, train_flag=True):
        self.imgs_info = self.get_images(txt_path)
        self.train_flag = train_flag

        self.train_tf = transforms.Compose([
            transforms.Resize(224),
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.ToTensor(),
        ])
        self.val_tf = transforms.Compose([
            transforms.Resize(224),
            transforms.ToTensor(),
        ])

    def get_images(self, txt_path):
        with open(txt_path, 'r', encoding='utf-8') as f:
            imgs_info = f.readlines()
            imgs_info = list(map(lambda x:x.strip().split('\t'), imgs_info))
        return imgs_info

    def padding_black(self, img):
        w, h = img.size
        scale = 224. / max(w, h)
        img_fg = img.resize([int(x) for x in [w*scale, h*scale]])
        size_fg = img_fg.size

        size_bg = 224
        img_bg = Image.new("RGB", (size_bg,size_bg))#(width,height)
        img_bg.paste(img_fg, ((size_bg - size_fg[0]) // 2,
                              (size_bg - size_fg[1] // 2))) # //表示整数除法
        img = img_bg
        return img

    def __getitem__(self, index):
        img_path, label = self.imgs_info[index] #获得第index个路径和标签
        img = Image.open(img_path)  #获得图片内容
        img = img.convert('RGB')
        img = self.padding_black(img)
        if self.train_flag:
            img = self.train_tf(img)
        else:
            img = self.val_tf(img)
        label = int(label)
        return img,label

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

if __name__ == "__main__":
    train_dataset = Garbage_Loader("train.txt", True)
    print("数据个数：",len(train_dataset))
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=1,
                                               shuffle=True)