Pytorch dataloader，线程太多，cpu内存分配太多

Question

我正在使用PyTorch训练一个模型。为了加载数据，我使用了torch.utils.data.DataLoader。数据加载器使用的是我实现的自定义数据库。出现了一个奇怪的问题，每次执行以下代码中的第二个for时，线程/进程的数量都会增加，并且会分配大量内存

    for epoch in range(start_epoch, opt.niter + opt.niter_decay + 1):
        epoch_start_time = time.time()
        if epoch != start_epoch:
            epoch_iter = epoch_iter % dataset_size
        for i, item in tqdm(enumerate(dataset, start=epoch_iter)):

我怀疑之前迭代器的线程和内存并没有在每次调用数据加载器的__iter__()之后被释放。分配的内存接近创建线程时由主线程/进程分配的内存量。也就是说，在初始时期，主线程使用2 2GB的内存，因此创建了两个大小为2 2GB的线程。在下一个时期，主线程将分配5 5GB的内存，并构造两个5 5GB的线程(num_workers为2)。我怀疑fork()函数会将大部分上下文复制到新线程中。

下面的活动监视器显示了python创建的进程，ZMQbg/1是与python相关的进程。

​

数据加载器使用的数据集有100个子数据集，__getitem__调用随机选择一个(忽略index)。(子数据集是来自pix2pixHD GitHub存储库的AlignedDataset )：

Answer 1

torch.utils.data.DataLoader预取2*num_workers，这样您就可以随时将数据发送到GPU/CPU，这可能是您看到内存增加的原因。

https://pytorch.org/docs/stable/_modules/torch/utils/data/dataloader.html