| | """ distributed launcher adopted from torch.distributed.launch
|
| | usage example: https://github.com/facebookresearch/maskrcnn-benchmark
|
| | This enables using multiprocessing for each spawned process (as they are treated as main processes)
|
| | """
|
| | import sys
|
| | import subprocess
|
| | from argparse import ArgumentParser, REMAINDER
|
| |
|
| | from utils import str2bool, int2str
|
| |
|
| | def parse_args():
|
| | parser = ArgumentParser(description="PyTorch distributed training launch "
|
| | "helper utilty that will spawn up "
|
| | "multiple distributed processes")
|
| |
|
| |
|
| | parser.add_argument('--n_GPUs', type=int, default=1, help='the number of GPUs for training')
|
| |
|
| |
|
| | parser.add_argument("training_script", type=str,
|
| | help="The full path to the single GPU training "
|
| | "program/script to be launched in parallel, "
|
| | "followed by all the arguments for the "
|
| | "training script")
|
| |
|
| |
|
| | parser.add_argument('training_script_args', nargs=REMAINDER)
|
| | return parser.parse_args()
|
| |
|
| | def main():
|
| | args = parse_args()
|
| |
|
| | processes = []
|
| | for rank in range(0, args.n_GPUs):
|
| | cmd = [sys.executable]
|
| |
|
| | cmd.append(args.training_script)
|
| | cmd.extend(args.training_script_args)
|
| |
|
| | cmd += ['--distributed', 'True']
|
| | cmd += ['--launched', 'True']
|
| | cmd += ['--n_GPUs', str(args.n_GPUs)]
|
| | cmd += ['--rank', str(rank)]
|
| |
|
| | process = subprocess.Popen(cmd)
|
| | processes.append(process)
|
| |
|
| | for process in processes:
|
| | process.wait()
|
| | if process.returncode != 0:
|
| | raise subprocess.CalledProcessError(returncode=process.returncode,
|
| | cmd=cmd)
|
| |
|
| | if __name__ == "__main__":
|
| | main()
|
| |
|
