Dataset+DataLoader加载数据流程代码详解

原创

已于 2023-07-20 16:09:14 修改 · 672 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能

于 2023-07-19 15:22:59 首次发布

引言

这篇文章是上一篇文章的一个补充，主要是从数据存储的角度出发去探索整个pytorch的数据加载方式，最终会追朔到数据文件的读写。本文以MNIST数据集举例，从源代码分析了解整个流程。

第一步定义数据集

train_data = torchvision.datasets.MNIST(
   root='./MINIST',  # 数据集的位置
   train=True,  # 如果为True则为训练集，如果为False则为测试集
   transform=torchvision.transforms.ToTensor(),  # 将图片转化成取值[0,1]的Tensor用于网络处理
   download=False
)

上面代码是初始化MNIST数据集的代码，MNIST是pytorch的torch.util.data.Dataset的实现。

class MNIST(VisionDataset):
class VisionDataset(data.Dataset):

可以看到它实现了Dataset中的中的__getitem__()接口，具体实现如下：

def __init__(
   self,
   root: str,
   train: bool = True,
   transform: Optional[Callable] = None,
   target_transform: Optional[Callable] = None,
   download: bool = False,
) -> None:
   super().__init__(root, transform=transform, target_transform=target_transform)
   self.train = train  # training set or test set
   if self._check_legacy_exist():
       self.data, self.targets = self._load_legacy_data()
       return
   if download:
       self.download()
   if not self._check_exists():
       raise RuntimeError("Dataset not found. You can use download=True to download it")
   self.data, self.targets = self._load_data()
def __getitem__(self, index: int) -> Tuple[Any, Any]:
   img, target = self.data[index], int(self.targets[index