【PyTorch】使用DataLoader自定义数据集读取

为了方便之后使用PyTorch的distributed部署&＃xff0c;加速训练&＃xff0c;将数据读取的方式改为适配pytorch提供的Dataset和DataLoader的方式。这里记录一下修改的要点&＃xff1a;

1. 涉及的import库&＃xff1a;

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


2. 自定义一个Dataset类&＃xff1a;

• 该类继承Dataset;

• 可以定义若干个数据预处理的函数&＃xff0c;关键的两个函数是&＃xff1a;__len__()和__getitem__();

• __getitem__()实际是python支持的一个迭代器函数&＃xff0c;编写时每次返回一个sample&＃xff0c;不需要定义batch size&＃xff0c;之后的DataLoader会自动帮忙读取数据组成batch的&＃xff1b;

• 举个栗子&＃xff1a;

  class MyDataset(Dataset):def __init__(self,data):self.data &＃61; datadef __len__(self):return len(self.data)def __getitem__(self):return self.datadef output(self):print(&＃39;output&＃39;)


3. 初始化Dataset和DataLoader类&＃xff1a;

• DataLoader的参数可参考&＃xff1a;https://blog.csdn.net/zyq12345678/article/details/90268668

• 注意&＃xff0c;如果在Dataset中每次返回的是自己定义的数据类型&＃xff0c;或者是字典类型&＃xff0c;有时要自己编写collate_fn()函数&＃xff0c;告诉系统如何返回一个batch。

• 举个栗子&＃xff1a;

  dataset &＃61; MyDataset(data)
dataloader &＃61; DataLoader(dataset,batch_size &＃61; 2,num_workers &＃61; 8,collate_fn &＃61; collate_fn,pin_memory &＃61; True
)
# 返回数据结构较复杂&＃xff0c;包括自定义数据类型或字典时
def collate_fn(batch):data &＃61; list(batch)return (data)


• 如果遇到类似报错&＃xff1a;

  TypeError: can&＃39;t pickle _thread._local objects

  请将DataLoader中的num_workers参数设置为0&＃xff0c;关闭多线程。原因可能是无法自动多线程处理复杂的数据类型。

4. 访问Dataloader内的Dataset类函数

• 举个栗子&＃xff1a;

for step, batch in enumerate(dataloader):dataloader.dataset.output()