Update train and test scripts for RPN optimizations

fschaeffler93 · Jul 9, 2019 · c175dfb · c175dfb
1 parent 413384a
commit c175dfb
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 0 deletions.
diff --git a/test.py b/test.py
@@ -18,6 +18,10 @@
 parser = argparse.ArgumentParser(description='testing')
 parser.add_argument('-n', '--tag', type=str, nargs='?', default='pre_trained_car',
                     help='set log tag')
+parser.add_argument('-d', '--decrease', type=bool, nargs='?', default=False,
+                    help='set the flag to True if decrease model')
+parser.add_argument('-m', '--minimize', type=bool, nargs='?', default=False,
+                    help='set the flag to True if minimize model')
 parser.add_argument('-b', '--single-batch-size', type=int, nargs='?', default=1,
                     help='set batch size for each gpu')
 parser.add_argument('-o', '--output-path', type=str, nargs='?', default='predictions',
@@ -56,6 +60,8 @@ def main(_):
         with tf.Session(config=config) as sess:
             model = RPN3D(
                 cls=cfg.DETECT_OBJ,
+                decrease=args.decrease,
+                minimize=args.minimize,
                 single_batch_size=args.single_batch_size,
                 avail_gpus=cfg.GPU_AVAILABLE.split(',')
             )

diff --git a/test_single.py b/test_single.py
@@ -18,6 +18,10 @@
 parser = argparse.ArgumentParser(description='testing')
 parser.add_argument('-n', '--tag', type=str, nargs='?', default='pre_trained_car',
                     help='set log tag')
+parser.add_argument('-d', '--decrease', type=bool, nargs='?', default=False,
+                    help='set the flag to True if decrease model')
+parser.add_argument('-m', '--minimize', type=bool, nargs='?', default=False,
+                    help='set the flag to True if minimize model')
 parser.add_argument('-t', '--data-tag', type=str, nargs='?', default='000000',
                     help='set data tag')
 parser.add_argument('-o', '--output-path', type=str, nargs='?',default='./predictions',
@@ -56,6 +60,8 @@ def main(_):
         with tf.Session(config=config) as sess:
             model = RPN3D(
                 cls=cfg.DETECT_OBJ,
+                decrease=args.decrease,
+                minimize=args.minimize,
                 single_batch_size=1,
                 avail_gpus=cfg.GPU_AVAILABLE.split(',')
             )

diff --git a/train.py b/train.py
@@ -22,6 +22,10 @@
                     help='max epoch')
 parser.add_argument('-n', '--tag', type=str, nargs='?', default='default',
                     help='set log tag')
+parser.add_argument('-d', '--decrease', type=bool, nargs='?', default=False,
+                    help='set the flag to True if decrease model')
+parser.add_argument('-m', '--minimize', type=bool, nargs='?', default=False,
+                    help='set the flag to True if minimize model')
 parser.add_argument('-b', '--single-batch-size', type=int, nargs='?', default=1,
                     help='set batch size for each gpu')
 parser.add_argument('-l', '--lr', type=float, nargs='?', default=0.001,
@@ -68,6 +72,8 @@ def main(_):
         with tf.Session(config=config) as sess:
             model = RPN3D(
                 cls=cfg.DETECT_OBJ,
+                decrease=args.decrease,
+                minimize=args.minimize,
                 single_batch_size=args.single_batch_size,
                 learning_rate=args.lr,
                 max_gradient_norm=5.0,