diff --git a/references/detection/train.py b/references/detection/train.py index 3b928611b4f..507d4faebae 100644 --- a/references/detection/train.py +++ b/references/detection/train.py @@ -8,6 +8,14 @@ The default hyperparameters are tuned for training on 8 gpus and 2 images per gpu. --lr 0.02 --batch-size 2 --world-size 8 If you use different number of gpus, the learning rate should be changed to 0.02/8*$NGPU. + +On top of that, for training Faster/Mask R-CNN, the default hyperparameters are + --epochs 26 --lr-steps 16 22 --aspect-ratio-group-factor 3 + +Also, if you train Keypoint R-CNN, the default hyperparameters are + --epochs 46 --lr-steps 36 43 --aspect-ratio-group-factor 3 +Because the number of images is smaller in the person keypoint subset of COCO, +the number of epochs should be adapted so that we have the same number of iterations. """ import datetime import os @@ -145,7 +153,7 @@ def main(args): parser.add_argument('--device', default='cuda', help='device') parser.add_argument('-b', '--batch-size', default=2, type=int, help='images per gpu, the total batch size is $NGPU x batch_size') - parser.add_argument('--epochs', default=13, type=int, metavar='N', + parser.add_argument('--epochs', default=26, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') @@ -158,12 +166,12 @@ def main(args): metavar='W', help='weight decay (default: 1e-4)', dest='weight_decay') parser.add_argument('--lr-step-size', default=8, type=int, help='decrease lr every step-size epochs') - parser.add_argument('--lr-steps', default=[8, 11], nargs='+', type=int, help='decrease lr every step-size epochs') + parser.add_argument('--lr-steps', default=[16, 22], nargs='+', type=int, help='decrease lr every step-size epochs') parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma') parser.add_argument('--print-freq', default=20, type=int, help='print frequency') parser.add_argument('--output-dir', default='.', help='path where to save') parser.add_argument('--resume', default='', help='resume from checkpoint') - parser.add_argument('--aspect-ratio-group-factor', default=0, type=int) + parser.add_argument('--aspect-ratio-group-factor', default=3, type=int) parser.add_argument( "--test-only", dest="test_only",