add benchmark set, reorder parameter of custom dataset

xvjiarui · Dec 21, 2018 · dc5edc3 · dc5edc3
1 parent f0ecb9d
commit dc5edc3
Show file tree

Hide file tree

Showing 6 changed files with 10 additions and 2 deletions.
diff --git a/configs/pascal_voc/ssd300_voc.py b/configs/pascal_voc/ssd300_voc.py
@@ -1,3 +1,4 @@
+benchmark = True
 # model settings
 input_size = 300
 model = dict(

diff --git a/configs/pascal_voc/ssd512_voc.py b/configs/pascal_voc/ssd512_voc.py
@@ -1,3 +1,4 @@
+benchmark = True
 # model settings
 input_size = 512
 model = dict(

diff --git a/configs/ssd300_coco.py b/configs/ssd300_coco.py
@@ -1,3 +1,4 @@
+benchmark = True
 # model settings
 input_size = 300
 model = dict(

diff --git a/configs/ssd512_coco.py b/configs/ssd512_coco.py
@@ -1,3 +1,4 @@
+benchmark = True
 # model settings
 input_size = 512
 model = dict(

diff --git a/mmdet/datasets/custom.py b/mmdet/datasets/custom.py
@@ -47,9 +47,9 @@ def __init__(self,
                  with_mask=True,
                  with_crowd=True,
                  with_label=True,
-                 test_mode=False,
                  extra_aug=None,
-                 resize_keep_ratio=True):
+                 resize_keep_ratio=True,
+                 test_mode=False):
         # prefix of images path
         self.img_prefix = img_prefix
 

diff --git a/tools/train.py b/tools/train.py
@@ -8,6 +8,7 @@
 from mmdet.apis import (train_detector, init_dist, get_root_logger,
                         set_random_seed)
 from mmdet.models import build_detector
+import torch
 
 
 def parse_args():
@@ -42,6 +43,9 @@ def main():
     args = parse_args()
 
     cfg = Config.fromfile(args.config)
+    # set benchmark
+    if cfg.get('benchmark', False):
+        torch.backends.cudnn.benchmark = True
     # update configs according to CLI args
     if args.work_dir is not None:
         cfg.work_dir = args.work_dir