For (Distributed)DataParallelism, main(args) spawn each process (main_worker) to each GPU.
(args)
| 20 | |
| 21 | |
| 22 | def main(args): |
| 23 | ''' |
| 24 | For (Distributed)DataParallelism, |
| 25 | main(args) spawn each process (main_worker) to each GPU. |
| 26 | ''' |
| 27 | |
| 28 | save_path = os.path.join(args.save_dir, args.save_name) |
| 29 | if os.path.exists(save_path) and args.overwrite and args.resume == False: |
| 30 | import shutil |
| 31 | shutil.rmtree(save_path) |
| 32 | if os.path.exists(save_path) and not args.overwrite: |
| 33 | raise Exception('already existing model: {}'.format(save_path)) |
| 34 | if args.resume: |
| 35 | if args.load_path is None: |
| 36 | raise Exception('Resume of training requires --load_path in the args') |
| 37 | if os.path.abspath(save_path) == os.path.abspath(args.load_path) and not args.overwrite: |
| 38 | raise Exception('Saving & Loading pathes are same. \ |
| 39 | If you want over-write, give --overwrite in the argument.') |
| 40 | |
| 41 | if args.seed is not None: |
| 42 | warnings.warn('You have chosen to seed training. ' |
| 43 | 'This will turn on the CUDNN deterministic setting, ' |
| 44 | 'which can slow down your training considerably! ' |
| 45 | 'You may see unexpected behavior when restarting ' |
| 46 | 'from checkpoints.') |
| 47 | |
| 48 | if args.gpu is not None: |
| 49 | warnings.warn('You have chosen a specific GPU. This will completely ' |
| 50 | 'disable data parallelism.') |
| 51 | |
| 52 | if args.dist_url == "env://" and args.world_size == -1: |
| 53 | args.world_size = int(os.environ["WORLD_SIZE"]) |
| 54 | |
| 55 | # distributed: true if manually selected or if world_size > 1 |
| 56 | args.distributed = args.world_size > 1 or args.multiprocessing_distributed |
| 57 | ngpus_per_node = torch.cuda.device_count() # number of gpus of each node |
| 58 | |
| 59 | if args.multiprocessing_distributed: |
| 60 | # now, args.world_size means num of total processes in all nodes |
| 61 | args.world_size = ngpus_per_node * args.world_size |
| 62 | |
| 63 | # args=(,) means the arguments of main_worker |
| 64 | mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args)) |
| 65 | else: |
| 66 | main_worker(args.gpu, ngpus_per_node, args) |
| 67 | |
| 68 | |
| 69 | def main_worker(gpu, ngpus_per_node, args): |
no test coverage detected