我正在尝试使用来自 2 个不同数据集的 2 个数据加载器来训练我的模型。
我找到了如何使用cycle() and zip()来设置它，因为我的数据集与这里的长度不同：如何使用 pytorch 同时迭代两个数据加载器？
File "/home/Desktop/example/train.py", line 229, in train_2
for i, (x1, x2) in enumerate(zip(cycle(train_loader_1), train_loader_2)):
File "/home/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/utils/data/dataloader.py", line 346, in __next__
data = self.dataset_fetcher.fetch(index) # may raise StopIteration
File "/home/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
return self.collate_fn(data)
File "/home/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py", line 80, in default_collate
return [default_collate(samples) for samples in transposed]
File "/home/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py", line 80, in <listcomp>
return [default_collate(samples) for samples in transposed]
File "/home/.conda/envs/3dcnn/lib/python3.7/site-packages/torch/utils/data/_utils/collate.py", line 56, in default_collate
return torch.stack(batch, 0, out=out)
RuntimeError: [enforce fail at CPUAllocator.cpp:64] . DefaultCPUAllocator: can't allocate memory: you tried to allocate 154140672 bytes. Error code 12 (Cannot allocate memory)
我试图通过设置num_workers=0来解决这个问题，减少批量大小，使用pinned_memory=False和shuffle=False...... 但这些都没有奏效...我有 256GB 的 RAM 和 4 个 NVIDIA TESLA V100 GPU。
我试图通过不同时而是单独在 2 个数据加载器中训练来运行它，它奏效了。但是对于我的项目，我需要使用 2 个数据集进行并行训练......
基于此讨论，我不是cycle()和zip()而是通过使用以下方法避免任何错误：
try:
data, target = next(dataloader_iterator)
except StopIteration:
dataloader_iterator = iter(dataloader)
data, target = next(dataloader_iterator)
从这个 PyTorch 帖子中@srossi93致敬！