Added --data-path to densenet and alexnet

Author: graham63

applications/vision/alexnet.py

@@ -22,6 +22,9 @@
 parser.add_argument(
     '--num-classes', action='store', default=1000, type=int,
     help='number of ImageNet classes (default: 1000)', metavar='NUM')
+parser.add_argument(
+    '--data-path', action='store', default=None, type=str,
+    help='Path to top-level imagenet directory. default: None')
 lbann.contrib.args.add_optimizer_arguments(parser)
 args = parser.parse_args()
 
@@ -64,7 +67,8 @@
 opt = lbann.contrib.args.create_optimizer(args)
 
 # Setup data reader
-data_reader = data.imagenet.make_data_reader(num_classes=args.num_classes)
+data_reader = data.imagenet.make_data_reader(num_classes=args.num_classes,
+                                             data_path=args.data_path)
 
 # Setup trainer
 trainer = lbann.Trainer(mini_batch_size=args.mini_batch_size)

applications/vision/densenet.py

@@ -336,6 +336,9 @@ def get_args():
     parser.add_argument(
         '--num-classes', action='store', default=1000, type=int,
         help='number of ImageNet classes (default: 1000)', metavar='NUM')
+    parser.add_argument(
+    '--data-path', action='store', default=None, type=str,
+    help='Path to top-level imagenet directory. default: None')
     lbann.contrib.args.add_optimizer_arguments(
         parser,
         default_optimizer='sgd',
@@ -379,7 +382,8 @@ def set_up_experiment(args,
                         callbacks=callbacks)
 
     # Set up data reader
-    data_reader = data.imagenet.make_data_reader(num_classes=args.num_classes)
+    data_reader = data.imagenet.make_data_reader(num_classes=args.num_classes,
+                                                 data_path=args.data_path)
 
     # Set up optimizer
     if args.optimizer == 'sgd':