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