NMR / tools /data_process /process /launch_process.py
Xxx999's picture
upload
45950ff
import os
from multiprocessing import Process
num_workers = 8
worker_ids = list(range(num_workers))
cuda_devices = [i % 8 for i in range(num_workers)]
def func(worker_id, cuda_device):
os.system(f"CUDA_VISIBLE_DEVICES={cuda_device} python -m data.process_seed_smpl --start_idx {worker_id} --interval {num_workers}")
if __name__ == "__main__":
works = []
for worker_id, cuda_device in zip(worker_ids, cuda_devices):
process = Process(target=func, args=(worker_id, cuda_device))
process.start()
works.append(process)
for process in works:
process.join()