()
| 104 | return args |
| 105 | |
| 106 | def main(): |
| 107 | args = get_args() |
| 108 | |
| 109 | if args.seed is not None or args.cudnn_deterministic: |
| 110 | seed_everything(args.seed, args.cudnn_deterministic) |
| 111 | |
| 112 | if args.gpu is not None: |
| 113 | warnings.warn('You have chosen a specific GPU. This will completely disable ddp.') |
| 114 | torch.cuda.set_device(args.gpu) |
| 115 | args.ngpus_per_node = 1 |
| 116 | args.world_size = 1 |
| 117 | else: |
| 118 | if args.num_node == 1: |
| 119 | args.dist_url == "auto" |
| 120 | else: |
| 121 | assert args.num_node > 1 |
| 122 | args.ngpus_per_node = torch.cuda.device_count() |
| 123 | args.world_size = args.ngpus_per_node * args.num_node |
| 124 | |
| 125 | launch(main_worker, args.ngpus_per_node, args.num_node, args.node_rank, args.dist_url, args=(args,)) |
| 126 | |
| 127 | |
| 128 | def main_worker(local_rank, args): |
no test coverage detected