add image_mobile.py

Author: implus

README.md

@@ -46,7 +46,7 @@ Single crop validation error on ImageNet-1k (center 224x224 crop from resized im
 #### ShuffleNetV2_1x
 ```
 python -m torch.distributed.launch --nproc_per_node=8 imagenet_mobile.py --cos -a shufflenetv2_1x --data /path/to/imagenet1k/ \
---epochs 300 --wd 4e-5 --gamma 0.1 -c checkpoints/imagenet/shufflenetv2_1x --train-batch 128 --opt-level O0 # Triaing
+--epochs 300 --wd 4e-5 --gamma 0.1 -c checkpoints/imagenet/shufflenetv2_1x --train-batch 128 --opt-level O0 --nowd-bn # Triaing
 
 python -m torch.distributed.launch --nproc_per_node=2 imagenet_mobile.py -a shufflenetv2_1x --data /path/to/imagenet1k/ \
 -e --resume ../pretrain/shufflenetv2_1x.pth.tar --test-batch 100 --opt-level O0 # Testing, ~69.6% top-1 Acc

classification/imagenet_fast.py

@@ -401,7 +401,8 @@ def main():
         print('==> Resuming from checkpoint..', args.resume)
         assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!'
         args.checkpoint = os.path.dirname(args.resume)
-        checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
+        #checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
+        checkpoint = torch.load(args.resume, map_location = lambda storage, loc: storage.cuda(args.gpu))
         best_acc = checkpoint['best_acc']
         start_epoch = checkpoint['epoch']
         # model may have more keys
@@ -529,7 +530,11 @@ def train(train_loader, model, criterion, optimizer, epoch, use_cuda):
         # loss.backward()
         with amp.scale_loss(old_loss, optimizer) as loss:
             loss.backward()
-        optimizer.step(print_flag=print_flag)
+
+        if args.el2:
+            optimizer.step(print_flag=print_flag)
+        else:
+            optimizer.step()
 
 
         if batch_idx % args.print_freq == 0: