使用PyTorch炼丹的过程中，我们最怕的就是在DataLoader里debug，原因无他：多进程驱动的DataLoader很难给出清晰的traceback报错，即便将num_worker设为0不启用多进程，有时一个DataLoader Worker PID XXX is killed by signal: Killed或者Segmentation Fault还是能让用户一脸懵逼。

新手炼丹师最常见到的问题就是在DataLoader的__getitem__方法调用cv2读取数据导致的segmentation fault，原因是cv2的多线程读取与DataLoader多进程读取交织会不定期造成死锁；或者是OS共享内存耗尽导致的DataLoader worker killed by signal。多年积赞的经验让笔者自以为在PyTorch调试技巧方面已经修炼的炉火纯青，可惜最近还是在一个困扰PyTorch社区3年的问题上栽了跟头。下面记录一下排查经过：

Update on 2023年03月04日：
这篇英文博客定量分析了各种multiprocessing方式（spawn/fork/forkserver）在不同设置下的内存泄露情况，如果你希望深入了解，还请拜读！如果你只希望定性了解解决方案，下面的内容也能满足你的需求。

太长不看版

  1. 不要在Python多进程setting下使用大的内建对象（list, dict, ...），而应使用有C binding的库提供的数据对象（如numpy.Array，pyarrow.Array）；
  2. 单机多实验进行data loading时，尝试读取存储在/dev/shm/中的memmap规避在每个实验中重复IO数据集带来的overhead。

起因

最近在单机开展多个实验时发现，实验数量一多，其中某个实验就会不定期报DataLoader Worker PID XXX is killed by signal: Killed；观察Memory Curve发现内存占用稳定增长，直到每个epoch结束后释放，一旦内存占满就报Killed：

image

排查

这其实是多进程编程中老生常谈的问题：在基于fork的多进程实现中，每次fork会让子进程得到不同的虚拟空间地址，但此时其映射的还是父进程的物理内存空间，可以让子进程高效率读取父进程的memory数据。一旦子进程有写操作就会触发操作系统的copy-on-write异常，系统会拷贝出另一块空间供子进程使用。

但实际上，除了数据增广，一般没有人会在DataLoader里改写对象吧？笔者上面的例子对大内存对象的唯一的操作是这样一句话：

scan = self.scans[ref['scan_id']]   # self.scans反序列化后是一个占用9GB内存的大dict

这句话是在一个继承Dataset类的数据集类的__init__方法里。从常识来看，这句话是对self.scans的Access操作，不应该导致copy-on-write。

但经这个issue的commect提醒，我们可以回顾一下PyObject在C中的结构：

[object.h]
/* PyObject_HEAD defines the initial segment of every PyObject. */
#define PyObject_HEAD 
    _PyObject_HEAD_EXTRA 
    int ob_refcnt;   // 引用计数
    struct _typeobject *ob_type;

#define PyObject_VAR_HEAD 
    PyObject_HEAD 
    int ob_size; /* Number of items in variable part */

每一次在fork后的子进程里对大型Python内建对象的引用，都会增加对象的ob_refcnt，变相write了这个对象，这个操作进而触发copy-on-write也就不奇怪了。

解决方法

Numpy等科学计算库从底层提供了数值、字符串等常用数据的handler np.Array，一个改写的例子如下：

改写前：

class traindataset(Dataset):
    def __init__(self):
          self.utt2fpath = {'utt1': './t1.wav', 'utt2': './t2.wav' ...}
          self.utts = list(self.utt2fpath.keys())
    def __getitem__(self, item):
          utt = self.utts[item]
          fpath = self.utt2fpath[utt]
          ...

改写后：

import numpy as np
class traindataset(Dataset):
    def __init__(self):
          self.utt2fpath = np.array([['utt1', './t1.wav'], ['utt2', './t2.wav'], ...])
    def __getitem__(self, item):
          utt = self.utt2fpath[item][0]
          fpath = self.utt2fpath[item][1]
          ...

相信不少人都习惯了在Dataset中使用python内建的list存储路径以便在__getitem__中读取，看了本文之后，你完全就可以用np.array避免不必要的memory overhead了。

注意！np.array只支持数值类型及字符串的C-binding，其他的python自建对象仍然以PyObject形式存在并通过pickle/unpickle的方式兼容numpy的逻辑。

如何判断？

>> np.array({'d': 'd'})
array({'d': 'd'}, dtype=object)  # dtype为object，代表numpy只是提供了PyObject的Wrapper，无法解决问题

这里也给其他研究者提一个醒，回到我遇到的scan = self.scans[ref['scan_id']] ，尽管有方法可以改写dict为两组np.array，但是这个dict的value是一个自定义的对象！

class ScannetDataset(object):
    def __init__(self, top_scan_dir, idx_to_semantic_cls_file,
                 instance_cls_to_semantic_cls_file, axis_alignment_info_file):
        self.top_scan_dir = top_scan_dir
        self.idx_to_semantic_cls_dict = read_dict(idx_to_semantic_cls_file)
        self.semantic_cls_to_idx_dict = invert_dictionary(self.idx_to_semantic_cls_dict)
        self.instance_cls_to_semantic_cls_dict = read_dict(instance_cls_to_semantic_cls_file)
        self.semantic_cls_to_instance_cls_dict = defaultdict(list)
        for k, v in self.instance_cls_to_semantic_cls_dict.items():
            self.semantic_cls_to_instance_cls_dict[v].append(k)
        self.scans_axis_alignment_matrices = read_dict(axis_alignment_info_file)

这就很难办了：经过搜寻，我没有发现任何能在多进程中共享Python自定义对象的方法，这也启示我们，在写新任务的IO逻辑时，尽量设计numpy-compatible的数据结构！

更进一步：单机多实验时避免每次实验重新加载数据集

不少人应该也有这样的需求：grid-search时需要在单机的不同卡上跑不同超参的实验，但通用的IO方案很容易出现IO瓶颈，我们希望在所有实验中共享读取进内存的数据。

第零步：修改/dev/shm挂载的分区大小

默认的/dev/shm大小可能不满足需求，通过以下命令修改，重启后失效：

sudo mount -o size=60000M  -o  remount  /dev/shm

第一步：将数据整理为numpy格式

图片、文本、音频都很好通过预处理方式统一为向量形式。

第二步：Cache，将数据用memmap形式载入/dev/shm

memory_map可以将存储在磁盘上的内容映射到内存，考虑到我们有多个独立的主进程和无数DataLoader Worker需要读取数据进行IO，memory_map是一个很好的solution。Numpy提供了np.lib.format.open_memmap函数可以直接读取np.save的二进制文件，十分方便。同时考虑到读取效率，直接将数据cache到/dev/shm这个直接映射到内存的目录。

如果只是读，没有对mmap中数据的写操作，只需要两行code：

  • cache: np.save(fn, array, allow_pickle=False)

    • fn是/dev/shm中某个文件的file handler，命名规则可以自定义。
    • array是需要存储的数据，同上，请确认这个array的dtype不是object，否则allow_pickle不能为False，意味着之后的np.load需要重新反序列化，每个进程还是要加载一次这个object。

  • 其他worker的读取: np.lib.format.open_memmap(filename, mode='r')

    • 返回值等同于np.load的结果，只不过内存开销极少，相当于创建了一个memory的view。

如果涉及到读写，可能需要用到下面这个库，提供了对mmap的锁机制。

sharearray.py中的cache方法提供了读写的功能，建议使用。但不推荐多个独立进程读写mmap，毕竟没有人帮你测试这个锁机制是否完善...

一个demo供参考

"""
Each time after restarting the OS:
sudo mount -o size=60000M  -o  remount  /dev/shm
"""
import os
import sys

sys.path.append(os.getcwd())

import argparse
from utils import share_array
import multiprocessing as mp
import numpy as np
from tqdm import tqdm
import time

SCENE_LEN = 1513
CACHE_PATH = None
CACHE_FILE_PREFIX = 'sharearray_'
KEYS = None


def _check(res_list):
    res_set = set()
    for k in KEYS:
        res_set.add(len([p for p in res_list if p.endswith(k + '.npy')]))
    return len(res_set) == 1 and list(res_set)[0] == SCENE_LEN


def get_workers(x):
    if x == -1:
        x = mp.cpu_count()
    return x


def get_npy_list(npy_path):
    res = [os.path.join(npy_path, p) for p in os.listdir(npy_path) if p.endswith('.npy')]
    _check(res)  # some simple check
    return res


def cache_array(scene_path):
    scene_name = os.path.basename(scene_path)[:-4]  # scene0685_01_clip_region_feat
    share_array.cache(scene_name, lambda: np.load(scene_path, allow_pickle=False),
                      shm_path=CACHE_PATH,
                      prefix=CACHE_FILE_PREFIX)


def clean_up_cache():
    del_list = [os.path.join(CACHE_PATH, p) for p in os.listdir(CACHE_PATH) if os.path.basename(p).startswith(CACHE_FILE_PREFIX)]
    for file in del_list:
        try:
            os.remove(file)
        except FileNotFoundError:
            print("Not Found: {}".format(file))
    print('Clean up Done!')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--offline-2d-feat', type=str,
                        default='/data/meta-ScanNet/split_feat/',
                        help='The disk path for offline 2D features')
    parser.add_argument('--cache-path', type=str, default='/dev/shm',
                        help='Cache target path, best be `/dev/shm` or other high-speed mount point (SSD, etc.)')
    parser.add_argument('--workers', type=int, default=-1)
    parser.add_argument('--debug', action='store_true', default=False)
    parser.add_argument('--keys', nargs='+', default=['clip_region_feat', 'clip_scaled_region_feat', 'obj_feat',
                                                      'obj_coord', 'camera_pose', 'instance_id'])
    args = parser.parse_args()
    num_workers = get_workers(args.workers)
    CACHE_PATH = args.cache_path
    KEYS = args.keys
    try:
        scenes_list = get_npy_list(args.offline_2d_feat)
        if not args.debug:
            with mp.Pool(num_workers) as p:
                _ = list(tqdm(p.imap(cache_array, scenes_list), total=len(scenes_list)))
                # automatic join & close
        else:
            for scene_path in tqdm(scenes_list):
                cache_array(scene_path)

        print("Done Caching! Happy Training! Ctrl+C to clean up cache.")

        while True:  # ugly infinite loop
            try:
                time.sleep(1e9)
            except KeyboardInterrupt:
                break
        print('Done!')
    finally:
        print('Cleaning up cache...')
        clean_up_cache()

参考

Possible to share in-memory data between 2 separate processes?

Shared-memory objects in multiprocessing

Last modification：March 4th, 2023 at 01:32 am
© 允许规范转载
If you think my article is useful to you, please feel free to appreciate