File size: 612 Bytes
45950ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | 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()
|