PyTorch 数据集通过基本 I/O 操作泄漏内存

Question

我正在迭代一个

torch.utils.data.DataLoader

及其关联的

torch.utils.data.Dataset

。我注意到在数据集的

__getitem__

方法中更改一行时，出现以下错误：

RuntimeError: DataLoader worker (pid 10666) is killed by signal: Bus error. It is possible that dataloader's workers are out of shared memory. Please try to raise your shared memory limit.
ERROR: Unexpected bus error encountered in worker. This might be caused by insufficient shared memory (shm).

__getitem__

以前长这样：

def __getitem__(self, idx):
    datafilename = os.path.join(self.root_dir, self.labelfile.iloc[idx,2])

    X = np.loadtxt(datafilename, delimiter=',', dtype=np.int32)

    X = torch.tensor(X, dtype=torch.float)
    return X

之后像这样：

def __getitem__(self, idx):
    datafilename = os.path.join(self.root_dir, self.labelfile.iloc[idx,2])

    with open(datafilename, 'r') as f:
        X = [int(x) for x in f.readline().split(',')]

    X = torch.tensor(X, dtype=torch.float)
    return X

如果有任何不同，我正在使用 VSCode 调试器运行。即使

num_workers=1

和我在两台不同的机器上尝试过同样的错误，这种行为仍然存在。我相信这不是由于硬件，而可能是内存泄漏。第二个版本也快了大约 7 倍，所以我更喜欢使用那个版本。

Answer 1

这不是内存泄漏。相反，它是 Python 的一个特性。在使用 map 或 list 等本机 Python 结构时，您可能会遇到此问题。有一个名为cstl的

NEW TOOL

（https://github.com/fuzihaofzh/cstl）。它封装了 C++ STL 容器来解决这个问题。它支持多种类型，包括 numpy 和 pytorch 不支持的嵌套 map、list 和 set。 这是一个简单的例子，展示了它是如何解决问题的：

from torch.utils.data import Dataset, DataLoader
import numpy as np
import torch
import copy
import sys
import cstl
from tqdm.auto import tqdm


class DataIter(Dataset):
    def __init__(self):
        cnt = 24000000
        self.cnt = cnt
        #self.data = np.array([x for x in range(cnt)]) # Good
        #self.data = [x for x in range(cnt)] #Leaky
        #self.data = cstl.MapIntInt({i : i for i in range(24000000)})# Good
        self.data = cstl.VecInt(range(24000000)) # Good

        
    def __len__(self):
        return self.cnt

    def __getitem__(self, idx):
        data = self.data[idx]
        data = np.array([int(data)], dtype=np.int64)
        return torch.tensor(data)

train_data = DataIter()
train_loader = DataLoader(train_data, batch_size=300,
                          shuffle=True,
                          drop_last=True,
                          pin_memory=False,
                          num_workers=18)

for i, item in tqdm(enumerate(train_loader)):
    torch.cuda.empty_cache()
    if i % 1000 == 0:
        print(i)