pytorch入门（三）

在前两篇文章中我们了解了pytorch的一些基础概念以及模型训练的整体框架。今天我们继续介绍pytorch的一些常用包。

传送门：

pytorch入门（一）

pytorch入门（二）

Dataset&DataLoader

对于数据处理上，我们希望样本数据集代码和模型训练代码分开，这样便于获得更强的可读性和模块化。因此，pytorch提供了两个数据元素：

from torch.utils.data import Dataset
from torch.utils.data import DataLoader

Dataset存储样本示例以及其对应的标签；DataLoader包装了数据集的可迭代对象，方便访问样本，可以认为是一个数据加载器。

Dataset是一种抽象类，所有继承它的子类都应该包含__len__方法和__getitem__方法，前者表示数据集大小，后者支持实例化后的样本整数索引。一般情况下，继承此类的对象也都包含初始化方法__init__.

下面举一个例子：

import os
import pandas as pd
from torchvision.io import read_image


class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform


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


    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

img_dir是存储数据(这个例子应该是图片)路径，annotations_file是存储图片对应标签的CSV文件。然后，方法__len__返回的是数据集所有图片的个数；__getitem__方法是在给定样本的索引后，加载样本数据以及其对应的标签。两个tranform是做一些必要的转换，这里先不用管。

在训练模型的时候，一般只取批量样本进行训练，因此，DataLoader实现这种“小批量”训练过程：

from torch.utils.data import DataLoader


DataLoader(