()
| 98 | return optimizer |
| 99 | |
| 100 | def main(): |
| 101 | args = parser.parse_args() |
| 102 | |
| 103 | if args.seed is not None: |
| 104 | random.seed(args.seed) |
| 105 | torch.manual_seed(args.seed) |
| 106 | cudnn.deterministic = True |
| 107 | warnings.warn('You have chosen to seed training. ' |
| 108 | 'This will turn on the CUDNN deterministic setting, ' |
| 109 | 'which can slow down your training considerably! ' |
| 110 | 'You may see unexpected behavior when restarting ' |
| 111 | 'from checkpoints.') |
| 112 | |
| 113 | if args.gpu is not None: |
| 114 | warnings.warn('You have chosen a specific GPU. This will completely ' |
| 115 | 'disable data parallelism.') |
| 116 | |
| 117 | if args.dist_url == "env://" and args.world_size == -1: |
| 118 | args.world_size = int(os.environ["WORLD_SIZE"]) |
| 119 | |
| 120 | args.distributed = args.world_size > 1 or args.multiprocessing_distributed |
| 121 | |
| 122 | ngpus_per_node = torch.cuda.device_count() |
| 123 | if args.multiprocessing_distributed: |
| 124 | # Since we have ngpus_per_node processes per node, the total world_size |
| 125 | # needs to be adjusted accordingly |
| 126 | args.world_size = ngpus_per_node * args.world_size |
| 127 | # Use torch.multiprocessing.spawn to launch distributed processes: the |
| 128 | # main_worker process function |
| 129 | mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args)) |
| 130 | else: |
| 131 | # Simply call main_worker function |
| 132 | main_worker(args.gpu, ngpus_per_node, args) |
| 133 | |
| 134 | |
| 135 |
no test coverage detected