| | |
|---|
| | | args = parse_args() |
|---|
| | | |
|---|
| | | # setup local gpu_id |
|---|
| | | os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id) |
|---|
| | | if args.ngpu > 0: |
|---|
| | | os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id) |
|---|
| | | |
|---|
| | | # DDP settings |
|---|
| | | if args.ngpu > 1: |
|---|
| | |
|---|
| | | |
|---|
| | | # re-compute batch size: when dataset type is small |
|---|
| | | if args.dataset_type == "small": |
|---|
| | | if args.batch_size is not None: |
|---|
| | | if args.batch_size is not None and args.ngpu > 0: |
|---|
| | | args.batch_size = args.batch_size * args.ngpu |
|---|
| | | if args.batch_bins is not None: |
|---|
| | | if args.batch_bins is not None and args.ngpu > 0: |
|---|
| | | args.batch_bins = args.batch_bins * args.ngpu |
|---|
| | | |
|---|
| | | main(args=args) |
|---|
