Python 多个 CPU、GPU 上的多处理

Python multiprocessing on multiple CPUs, GPUs

我有 8 个 GPU,64 个 CPU 核 (multiprocessing.cpu_count()=64)

我正在尝试使用深度学习模型推断多个视频文件。我希望在 8 个 GPU 中的每一个上处理一些文件。对于每个 GPU,我想要使用不同的 6 CPU 个内核。

python 文件名下方:inference_{gpu_id}.py

Input1: GPU_id

Input2: Files to process for GPU_id

from torch.multiprocessing import Pool, Process, set_start_method
try:
     set_start_method('spawn', force=True)
except RuntimeError:
    pass

model = load_model(device='cuda:' + gpu_id) 

def pooling_func(file):
    preds = []
    cap = cv2.VideoCapture(file)
    while(cap.isOpened()):
          ret, frame = cap.read()
          count += 1
          if ret == True:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                pred = model(frame)[0]
                preds.append(pred)
          else:
                break
    cap.release()
    np.save(file[:-4]+'.npy', preds)

def process_files():

    # all files to process on gpu_id
    files = np.load(gpu_id + '_files.npy') 

    # I am hoping to use 6 cores for this gpu_id, 
    # and a different 6 cores for a different GPU id
    pool = Pool(6) 

    r = list(tqdm(pool.imap(pooling_func, files), total = len(files)))
    pool.close()
    pool.join()

if __name__ == '__main__':
    import multiprocessing
    multiprocessing.freeze_support()
    process_files()

我希望在所有 GPU 上同时 运行 inference_{gpu_id}.py 文件

目前,我能够在一个 6 核 GPU 上成功 运行 它,但是当我尝试在所有 GPU 上同时 运行 它时,只有 GPU 0 运行s , 所有其他人停止给出以下错误消息。

RuntimeError: CUDA error: invalid device ordinal.

脚本我是运行ning:

CUDA_VISIBLE_DEVICES=0 inference_0.py

CUDA_VISIBLE_DEVICES=1 inference_1.py

...

CUDA_VISIBLE_DEVICES=7 inference_7.py

以下原是对一个question you asked的回答,后来被删除了。


考虑一下,如果您没有使用 CUDA_VISIBLE_DEVICES 标志,那么所有 GPU 都将可用于您的 PyTorch 进程。这意味着 torch.cuda.device_count 将 return 8 (假设您的版本设置有效)。您将能够通过 torch.devicetorch.device('cuda:0')torch.device('cuda:1')、...和 ​​torch.device('cuda:8').[=37= 访问这 8 个 GPU 中的每一个]

现在,如果您只打算使用一个并且希望将您的进程限制为一个。然后 CUDA_VISIBLE_DEVICES=i (其中 i 是设备序号)将使它成为现实。在这种情况下,torch.cuda 将只能通过 torch.device('cuda:0') 访问单个设备。实际设备序号是什么并不重要,您访问它的方式是通过 torch.device('cuda:0').

如果您允许访问多个设备:假设 n°0、n°4 和 n°2,那么您将使用 CUDA_VISIBLE_DEVICES=0,4,2。因此,您通过 d0 = torch.device('cuda:0')d1 = torch.device('cuda:1')d2 = torch.device('cuda:2') 引用您的 cuda 设备。与您使用标志定义它们的顺序相同,:

d0 -> GPU n°0, d1 -> GPU n°4, and d2 -> GPU n°2.

这使得您可以在不同的 GPU 上使用相同的代码 运行 而无需更改底层代码 您指的是设备序号.

总而言之,您需要查看的是 运行 您的代码所需的设备数量。在您的情况下:1 就足够了。您将使用 torch.device('cuda:0') 来引用它。但是,当 运行ning 您的代码时,您将需要指定 cuda:0 设备是什么,并带有标志:

> CUDA_VISIBLE_DEVICES=0 inference.py
> CUDA_VISIBLE_DEVICES=1 inference.py
  ...
> CUDA_VISIBLE_DEVICES=7 inference.py

请注意 'cuda' 将默认为 'cuda:0'