Merge pull request #16 from developer0hye/patch-1

Cleanup

lr_finder.py

@@ -48,6 +48,7 @@ def get_lr(self):
     parser.add_argument('-base_lr', type=float, default=1e-7, help='min learning rate')
     parser.add_argument('-max_lr', type=float, default=10, help='max learning rate')
     parser.add_argument('-num_iter', type=int, default=100, help='num of iteration')
+    parser.add_argument('-gpu', type=bool, default=True, help='use gpu or not')
     parser.add_argument('-gpus', nargs='+', type=int, default=0, help='gpu device')
     args = parser.parse_args()
 

train.py

@@ -118,8 +118,8 @@ def eval_training(epoch):
     parser.add_argument('-lr', type=float, default=0.1, help='initial learning rate')
     args = parser.parse_args()
 
-    net = get_network(args, use_gpu=args.gpu)
-
+    net = get_network(args)
+        
     #data preprocessing:
     cifar100_training_loader = get_training_dataloader(
         settings.CIFAR100_TRAIN_MEAN,

utils.py

@@ -14,7 +14,7 @@
 from torch.utils.data import DataLoader
 
 
-def get_network(args, use_gpu=True):
+def get_network(args):
     """ return given network
     """
 
@@ -140,7 +140,7 @@ def get_network(args, use_gpu=True):
         print('the network name you have entered is not supported yet')
         sys.exit()
 
-    if use_gpu:
+    if args.gpu: #use_gpu
         net = net.cuda()
 
     return net
@@ -229,4 +229,4 @@ def get_lr(self):
         """we will use the first m batches, and set the learning
         rate to base_lr * m / total_iters
         """
-        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]
+        return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]