Merge pull request #17 from scottclowe/api_batchsize

API: Change main_lincls --batch-size argument to match main_moco
main
Xinlei Chen 2021-11-25 01:33:37 -08:00 committed by GitHub
commit c349e6e24f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 3 deletions

View File

@ -53,7 +53,7 @@ parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
parser.add_argument('-b', '--batch-size', default=1024, type=int,
metavar='N',
help='mini-batch size (default: 1024), this is the total '
'batch size of all GPUs on the current node when '
'batch size of all GPUs on all nodes when '
'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
metavar='LR', help='initial (base) learning rate', dest='lr')
@ -207,7 +207,7 @@ def main_worker(gpu, ngpus_per_node, args):
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.batch_size = int(args.batch_size / args.world_size)
args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
else:

View File

@ -60,7 +60,7 @@ parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
parser.add_argument('-b', '--batch-size', default=4096, type=int,
metavar='N',
help='mini-batch size (default: 4096), this is the total '
'batch size of all GPUs on the current node when '
'batch size of all GPUs on all nodes when '
'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr', '--learning-rate', default=0.6, type=float,
metavar='LR', help='initial (base) learning rate', dest='lr')