Update default, add note about scheduler

Author: 1adrianb

examples/imagenet.py

@@ -21,6 +21,7 @@
 
 from utils import *
 
+# install warmup_scheduler from https://github.com/ildoonet/pytorch-gradual-warmup-lr
 from warmup_scheduler import GradualWarmupScheduler
 import bnn.models.resnet as models
 from bnn.engine import BinaryChef
@@ -48,11 +49,11 @@
                     help='mini-batch size (default: 256), this is the total '
                          'batch size of all GPUs on the current node when '
                          'using Data Parallel or Distributed Data Parallel')
-parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
+parser.add_argument('--lr', '--learning-rate', default=0.001, type=float,
                     metavar='LR', help='initial learning rate', dest='lr')
 parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                     help='momentum')
-parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
+parser.add_argument('--wd', '--weight-decay', default=1e-2, type=float,
                     metavar='W', help='weight decay (default: 1e-4)',
                     dest='weight_decay')
 parser.add_argument('-p', '--print-freq', default=100, type=int,
@@ -190,7 +191,7 @@ def main_worker(gpu, ngpus_per_node, args):
 
     parameters = model.parameters()
     if args.optimizer == 'adamw':
-        wd = 0.01 if args.step == 0 else 0.01
+        wd = args.weight_decay if args.step == 0 else 0
         optimizer = torch.optim.AdamW(parameters, args.lr, weight_decay=wd)
     elif args.optimizer == 'adam':
         optimizer = torch.optim.Adam(parameters, args.lr)