PyTorch 1.11.0自定义数据集DataLoader随机冻结问题求助

使用PyTorch 1.11.0基于自定义数据集实现的DataLoader存在随机冻结问题，无法稳定复现：多数情况下运行无异常，但偶尔会陷入停滞。手动中断（ctrl+c）后得到如下报错栈：

idx, data = self._get_data()
  File "/opt/conda/envs/torch/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1163, in _get_data
    success, data = self._try_get_data()
  File "/opt/conda/envs/torch/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1011, in _try_get_data
    data = self._data_queue.get(timeout=timeout)
  File "/opt/conda/envs/torch/lib/python3.8/queue.py", line 179, in get
    self.not_empty.wait(remaining)
  File "/opt/conda/envs/torch/lib/python3.8/threading.py", line 306, in wait
    gotit = waiter.acquire(True, timeout)
KeyboardInterrupt

环境与前期排查 #

• 运行环境：Kubernetes上的Ubuntu 20.04 Linux Pod
• 单独遍历数据集无任何异常，仅遍历DataLoader时会出现随机停滞
• 已尝试的无效方案：
  • 将pin_memory设为True，问题依旧
  • 将num_workers设为0可避免冻结，但速度过慢；设为任何大于0的值仍会出现冻结
• 官方论坛多次提问未获回复，现寻求排查方向

自定义数据集代码 #

class MultiModalDataset(Dataset):
    def __init__(self, img_dataset: pd.DataFrame, text_dataset: pd.DataFrame, 
            img_fld: str, img_transforms=None, n_classes=None, img_size=224,
            n_sentences=1, n_tokens=12, collate_fn=None, l1normalization=False, verbose=False):
        super().__init__()
        self.n_classes = n_classes or img_dataset.shape[1]
        assert self.n_classes == img_dataset.shape[1]
        self.img_ds = img_dataset
        # print(text_dataset.head())
        self.text_ds = text_dataset.set_index("image_filename")
        self.img_fld = img_fld
        self.transforms = img_transforms
        self.img_size = img_size
        self.n_sentences = n_sentences
        self.n_tokens = n_tokens
        self.collate_fn = collate_fn
        self.l1normalization = l1normalization
        self.verbose = verbose

    def __len__(self):
        return len(self.img_ds)
    
    def __getitem__(self, idx):
        assert (idx >=0) and (idx < len(self.img_ds))
        item = self.img_ds.iloc[idx]
        filename = item.name
        labels = item.values
        if self.l1normalization:
            nlabs = sum(labels)
            assert nlabs > 0, f"dataset, at index {idx}, no labels found"
            labels = labels / nlabs
        
        text = self.text_ds.loc[filename, "enc_text"]

        if self.collate_fn is not None:
            padded_text = self.collate_fn(text, n_sents=self.n_sentences, max_tokens=self.n_tokens, verbose=self.verbose)
        else:
            padded_text = text

        return self.load_image(filename), torch.tensor(labels.astype(np.float32)), torch.tensor(padded_text)

    def load_image(self, img_filename):
        fn = join(self.img_fld, img_filename)
        img = Image.open(fn)
        if self.transforms is not None:
            img = self.transforms(img)
        return img

DataLoader配置 #

DataLoader(dataset, batch_size=128, shuffle=True, num_workers=4, drop_last=[False,False,False], pin_memory=False)

遍历方式 #

for bi, (_, _, _) in enumerate(dataloader):
 ...

希望能得到针对该随机冻结问题的排查思路，尤其是结合Python并发机制和Kubernetes环境的可能原因。

内容的提问来源于stack exchange，提问作者Antonio Sesto