added batch_size argument to train cli

ArneSchulzTUBS · Nov 16, 2020 · 1340e6f · 1340e6f
1 parent 9b3cf80
commit 1340e6f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/orthophoto-segmentation-benchmark-toolkit/main_training.py b/orthophoto-segmentation-benchmark-toolkit/main_training.py
@@ -11,6 +11,7 @@
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument("--epochs", help="model weights of given eperiment will be used for training", default=30)
+    parser.add_argument("--bs", help="model weights of given eperiment will be used for training", default=4)
     parser.add_argument("--experiment", help="model weights of given eperiment will be used for training", default="")
     args = parser.parse_args()
 
@@ -21,7 +22,7 @@
     dataset = DroneDeployDataset(dataset, size).download().generate_chips()
     model_backend = UnetBackend('efficientnetb3')
 
-    experiment = Experiment("test", dataset, model_backend, batch_size=1,
+    experiment = Experiment("test", dataset, model_backend, batch_size=args.bs,
                             experiment_directory=args.experiment, load_best=False)
     experiment.analyze()
     experiment.train(epochs=args.epochs)