多卡多进程并发加速

import torch

import sys
import os
import multiprocessing as mp
from multiprocessing import Manager

os.chdir(os.path.dirname(os.path.abspath(__file__)))

current_dir = os.path.dirname(os.path.abspath(__file__))
paths = [
    *[os.path.abspath(os.path.join(current_dir, subdir)) for subdir in ['../','../../../../giga-datasets', '../../../../giga-models', '../../../../giga-train']]
]

sys.path[:0] = paths
os.environ['PYTHONPATH'] = ':'.join(paths + os.environ.get('PYTHONPATH', '').split(':'))


    mp.set_start_method('spawn')
    n_gpus=[1,2]
    # Create processes for each GPU
    processes = []
    for data_i, rank in  enumerate(n_gpus):
        p = mp.Process(target=val_dataset, args=(rank, config_path[data_i]))
        p.start()
        processes.append(p)

    # Wait for all processes to complete
    for p in processes:
        p.join()