Pytorch使用DataLoader, num_workers!=0时的内存泄露

PyTorch内存泄露：大数据集多分类训练中的问题与解决方案

原创

已于 2023-10-08 09:05:03 修改 · 2.5k 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python #DataLoader

于 2023-10-08 01:00:56 首次发布

描述一下背景，和遇到的问题：

我在做一个超大数据集的多分类，设备Ubuntu 22.04+i9 13900K+Nvidia 4090+64GB RAM，第一次的训练的训练集有700万张，训练成功。后面收集到更多数据集，数据增强后达到了1000万张。但第二次训练4个小时后，就被系统杀掉进程了，原因是Out of Memory。找了很久的原因，发现内存随着训练step的增加而线性增加，猜测是内存泄露，最后定位到了DataLoader的num_workers参数（只要num_workers=0就没有问题）。

真正原因：

Python（Pytorch）中的list转换成tensor时，会发生内存泄漏，要避免list的使用，可以通过使用np.array来代替list。

解决办法：

自定义DataLoader中的Dataset类，然后Dataset类中的list全部用np.array来代替。这样的话，DataLoader将np.array转换成Tensor的过程就不会发生内存泄露。

下面给两个错误的示例代码和一个正确的代码：（都是我自己犯过的错误）

1.错误的DataLoader加载数据集方法1

# 加载数据
train_data = datasets.ImageFolder(root=TRAIN_DIR_ARG, transform=transform)
valid_data = datasets.ImageFolder(root=VALIDATION_DIR, transform=transform)
test_data = datasets.ImageFolder(root=TEST_DIR, transform=transform)

train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True,