Added --data-path to densenet and alexnet

LLNL · Mar 6, 2023 · 5fbe41e · 5fbe41e
1 parent 7392542
commit 5fbe41e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/applications/vision/alexnet.py b/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)

diff --git a/applications/vision/densenet.py b/applications/vision/densenet.py
@@ -428,6 +428,8 @@ def get_args():
     parser.add_argument("--print-matrix-summary", dest="print_matrix_summary",
                         action="store_const",
                         const=True, default=False)
+    parser.add_argument('--data-path', action='store', default=None, type=str,
+                        help='Path to top-level imagenet directory. default: None')
     args = parser.parse_args()
     return args
 
@@ -438,7 +440,7 @@ def set_up_experiment(args,
                       labels):
     algo = lbann.BatchedIterativeOptimizer("sgd", epoch_count=args.num_epochs)
 
-    
+
     # Set up objective function
     cross_entropy = lbann.CrossEntropy([probs, labels])
     layers = list(lbann.traverse_layer_graph(input_))
@@ -472,7 +474,9 @@ def set_up_experiment(args,
                         callbacks=callbacks)
 
     # Set up data reader
-    data_reader = data.imagenet.make_data_reader(num_classes=args.num_classes, small_testing=True)
+    data_reader = data.imagenet.make_data_reader(num_classes=args.num_classes,
+                                                 small_testing=True,
+                                                 data_path=args.data_path)
 
     percentage = 0.001 * 2 * (args.mini_batch_size / 16) * 2