我正在 Pytorch 中训练图像分类模型,并使用它们的默认数据加载器来加载我的训练数据。我有一个非常大的训练数据集,所以通常每个班级有几千个样本图像。过去我训练过的模型总共有大约 20 万张图像,没有任何问题。但是我发现当总共有超过一百万张图像时,Pytorch 数据加载器会卡住。
我相信当我打电话时代码正在挂起datasets.ImageFolder(...)
。当我 Ctrl-C 时,这始终是输出:
Traceback (most recent call last): │
File "main.py", line 412, in <module> │
main() │
File "main.py", line 122, in main │
run_training(args.group, args.num_classes) │
File "main.py", line 203, in run_training │
train_loader = create_dataloader(traindir, tfm.train_trans, shuffle=True) │
File "main.py", line 236, in create_dataloader │
dataset = datasets.ImageFolder(directory, trans) │
File "/home/username/.local/lib/python3.5/site-packages/torchvision/datasets/folder.py", line 209, in __init__ │
is_valid_file=is_valid_file) │
File "/home/username/.local/lib/python3.5/site-packages/torchvision/datasets/folder.py", line 94, in __init__ │
samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file) │
File "/home/username/.local/lib/python3.5/site-packages/torchvision/datasets/folder.py", line 47, in make_dataset │
for root, _, fnames in sorted(os.walk(d)): │
File "/usr/lib/python3.5/os.py", line 380, in walk │
is_dir = entry.is_dir() │
Keyboard Interrupt
我认为某处可能存在死锁,但是根据 Ctrl-C 的堆栈输出,它看起来不像在等待锁定。所以后来我认为数据加载器很慢,因为我试图加载更多数据。我让它运行了大约 2 天,但没有任何进展,在加载的最后 2 小时内,我检查了 RAM 使用量保持不变。在过去不到几个小时的时间内,我还能够加载包含超过 20 万张图像的训练数据集。我还尝试将我的 GCP 机器升级为拥有 32 个内核、4 个 GPU 和超过 100GB 的 RAM,但似乎在加载了一定数量的内存后,数据加载器就会卡住。
我很困惑数据加载器在遍历目录时如何卡住,我仍然不确定它是卡住还是非常慢。有什么方法可以改变 Pytortch 数据加载器,使其能够处理超过 100 万张图像进行训练?任何调试建议也值得赞赏!
谢谢!