自定义 DataLoader 时应使用 Unix 系统
自定义 Dataset 类
PyTorch 允许自定义 Dataset 类,并由此获得 DataLoader,能方便训练时获得 batch:
from torch.utils.data import DataLoader, Dataset
import h5py
import os
class RadarDataset(Dataset):
def __init__(self, directory):
...
def __len__(self):
...
def __getitem__(self, idx):
...
dataset = RadarDataset(r'/mnt/z/automotive_pre_processed')
dataloader = DataLoader(
dataset,
batch_size=128,
shuffle=True,
num_workers=8,
persistent_workers=True,
prefetch_factor=3,
)
顺便一提,若要在自定义 Dataset 类中使用 h5py 库,要注意文件句柄不能在
__init__
里获取和存储。因为__init__
中出现的成员会被其他 worker 共享,而 h5py 不允许。
现在问题出现了。使用这个 dataloader
,总是出现错误:
dataloader worker (pid(s) 9144, 29312, 25764, 26220, 27448, 27116) exited unexpectedly
问题的解决方法
PyTorch 文档 提到过这个问题。
在Unix上,默认 fork() 启动多进程,子 worker 可以直接访问数据集和 Python 参数函数;
在 Windows 或 MacOS 上,默认 spawn() 启动多进程,即启动另一个解释器运行主脚本,然后内部的工作函数通过 pickle 序列化接收数据集、collate_fn和其他参数。
简单来说,用 Linux 吧,就没这些破事了。