|
""" 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()
|
|
|
