Fix retinahead

WallZFE · Dec 16, 2019 · de527fc · de527fc
1 parent d6e5b02
commit de527fc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/models/retinahead.py b/models/retinahead.py
@@ -57,7 +57,6 @@ def __init__(self,
             [2**(i / scales_per_octave) for i in range(scales_per_octave)])
         anchor_scales = octave_scales * octave_base_scale
         self.cls_out_channels = num_classes
-        print('num_class: ', self.cls_out_channels)
         self.num_anchors = len(self.anchor_ratios) * len(self.anchor_scales)
         self._init_layers()
     def _init_layers(self):

diff --git a/train.py b/train.py
@@ -32,9 +32,9 @@
                     help='Batch size for training')
 parser.add_argument('--num_worker', default=8, type=int,
                     help='Number of workers used in dataloading')
-parser.add_argument('--num_classes', default=80, type=int,
+parser.add_argument('--num_classes', default=20, type=int,
                     help='Number of class used in model')
-parser.add_argument('--device', default=[0], type=list,
+parser.add_argument('--device', default=[0, 1], type=list,
                     help='Use CUDA to train model')
 parser.add_argument('--grad_accumulation_steps', default=1, type=int,
                     help='Number of gradient accumulation steps')