yjh0410 1 year ago
parent
commit
37a622a25f

+ 1 - 1
odlab/models/backbone/__init__.py

@@ -3,7 +3,7 @@ from .resnet           import build_resnet
 
 
 def build_backbone(cfg):
 def build_backbone(cfg):
     print('==============================')
     print('==============================')
-    print('Backbone: {}'.format(cfg['backbone']))
+    print('Backbone: {}'.format(cfg.backbone))
     # ResNet
     # ResNet
     if "resnet" in cfg.backbone:
     if "resnet" in cfg.backbone:
         return build_resnet(cfg)
         return build_resnet(cfg)

+ 1 - 1
odlab/models/detectors/__init__.py

@@ -23,7 +23,7 @@ def build_model(args, cfg, is_val=False):
     if is_val:
     if is_val:
         # ------------ Keep training from the given weight ------------
         # ------------ Keep training from the given weight ------------
         if args.resume is not None:
         if args.resume is not None:
-            print('keep training: ', args.resume)
+            print('Load model from the checkpoint: ', args.resume)
             checkpoint = torch.load(args.resume, map_location='cpu')
             checkpoint = torch.load(args.resume, map_location='cpu')
             # checkpoint state dict
             # checkpoint state dict
             checkpoint_state_dict = checkpoint.pop("model")
             checkpoint_state_dict = checkpoint.pop("model")

+ 11 - 11
odlab/models/head/__init__.py

@@ -5,27 +5,27 @@ from .fcos_head      import FcosHead
 # build head
 # build head
 def build_head(cfg, in_dim, out_dim, num_classes):
 def build_head(cfg, in_dim, out_dim, num_classes):
     print('==============================')
     print('==============================')
-    print('Head: {}'.format(cfg['head']))
+    print('Head: {}'.format(cfg.head))
     
     
-    if cfg['head'] == 'fcos_head':
+    if cfg.head == 'fcos_head':
         model = FcosHead(cfg          = cfg,
         model = FcosHead(cfg          = cfg,
                          in_dim       = in_dim,
                          in_dim       = in_dim,
                          out_dim      = out_dim,
                          out_dim      = out_dim,
                          num_classes  = num_classes,
                          num_classes  = num_classes,
-                         num_cls_head = cfg['num_cls_head'],
-                         num_reg_head = cfg['num_reg_head'],
-                         act_type     = cfg['head_act'],
-                         norm_type    = cfg['head_norm']
+                         num_cls_head = cfg.num_cls_head,
+                         num_reg_head = cfg.num_reg_head,
+                         act_type     = cfg.head_act,
+                         norm_type    = cfg.head_norm,
                          )
                          )
-    elif cfg['head'] == 'yolof_head':
+    elif cfg.head == 'yolof_head':
         model = YolofHead(cfg          = cfg,
         model = YolofHead(cfg          = cfg,
                           in_dim       = in_dim,
                           in_dim       = in_dim,
                           out_dim      = out_dim,
                           out_dim      = out_dim,
                           num_classes  = num_classes,
                           num_classes  = num_classes,
-                          num_cls_head = cfg['num_cls_head'],
-                          num_reg_head = cfg['num_reg_head'],
-                          act_type     = cfg['head_act'],
-                          norm_type    = cfg['head_norm']
+                          num_cls_head = cfg.num_cls_head,
+                          num_reg_head = cfg.num_reg_head,
+                          act_type     = cfg.head_act,
+                          norm_type    = cfg.head_norm,
                           )
                           )
 
 
     return model
     return model

+ 1 - 1
odlab/models/head/fcos_head.py

@@ -37,7 +37,7 @@ class FcosHead(nn.Module):
         self.num_reg_head = num_reg_head
         self.num_reg_head = num_reg_head
         self.act_type = act_type
         self.act_type = act_type
         self.norm_type = norm_type
         self.norm_type = norm_type
-        self.stride = cfg['out_stride']
+        self.stride = cfg.out_stride
 
 
         # ------------------ Network parameters -------------------
         # ------------------ Network parameters -------------------
         ## cls head
         ## cls head

+ 4 - 4
odlab/models/head/yolof_head.py

@@ -9,7 +9,7 @@ class YolofHead(nn.Module):
     def __init__(self, cfg, in_dim, out_dim, num_classes, num_cls_head=1, num_reg_head=1, act_type='relu', norm_type='BN'):
     def __init__(self, cfg, in_dim, out_dim, num_classes, num_cls_head=1, num_reg_head=1, act_type='relu', norm_type='BN'):
         super().__init__()
         super().__init__()
         self.fmp_size = None
         self.fmp_size = None
-        self.ctr_clamp = cfg['center_clamp']
+        self.ctr_clamp = cfg.center_clamp
         self.DEFAULT_EXP_CLAMP = math.log(1e8)
         self.DEFAULT_EXP_CLAMP = math.log(1e8)
         self.DEFAULT_SCALE_CLAMP = math.log(1000.0 / 16)
         self.DEFAULT_SCALE_CLAMP = math.log(1000.0 / 16)
         # ------------------ Basic parameters -------------------
         # ------------------ Basic parameters -------------------
@@ -20,10 +20,10 @@ class YolofHead(nn.Module):
         self.num_reg_head=num_reg_head
         self.num_reg_head=num_reg_head
         self.act_type=act_type
         self.act_type=act_type
         self.norm_type=norm_type
         self.norm_type=norm_type
-        self.stride = cfg['out_stride']
+        self.stride = cfg.out_stride
         # Anchor config
         # Anchor config
-        self.anchor_size = torch.as_tensor(cfg['anchor_size'])
-        self.num_anchors = len(cfg['anchor_size'])
+        self.anchor_size = torch.as_tensor(cfg.anchor_size)
+        self.num_anchors = len(cfg.anchor_size)
 
 
         # ------------------ Network parameters -------------------
         # ------------------ Network parameters -------------------
         ## cls head
         ## cls head

+ 16 - 33
odlab/models/neck/__init__.py

@@ -1,5 +1,4 @@
 from .dilated_encoder import DilatedEncoder
 from .dilated_encoder import DilatedEncoder
-from .hybrid_encoder import HybridEncoder
 from .fpn import BasicFPN
 from .fpn import BasicFPN
 from .spp import SPPF
 from .spp import SPPF
 
 
@@ -7,51 +6,35 @@ from .spp import SPPF
 # build neck
 # build neck
 def build_neck(cfg, in_dim, out_dim):
 def build_neck(cfg, in_dim, out_dim):
     print('==============================')
     print('==============================')
-    print('Neck: {}'.format(cfg['neck']))
+    print('Neck: {}'.format(cfg.neck))
 
 
     # ----------------------- Neck module -----------------------
     # ----------------------- Neck module -----------------------
-    if cfg['neck'] == 'dilated_encoder':
+    if cfg.neck == 'dilated_encoder':
         model = DilatedEncoder(in_dim       = in_dim,
         model = DilatedEncoder(in_dim       = in_dim,
                                out_dim      = out_dim,
                                out_dim      = out_dim,
-                               expand_ratio = cfg['neck_expand_ratio'],
-                               dilations    = cfg['neck_dilations'],
-                               act_type     = cfg['neck_act'],
-                               norm_type    = cfg['neck_norm']
+                               expand_ratio = cfg.neck_expand_ratio,
+                               dilations    = cfg.neck_dilations,
+                               act_type     = cfg.neck_act,
+                               norm_type    = cfg.neck_norm,
                                )
                                )
-    elif cfg['neck'] == 'spp_block':
+    elif cfg.neck == 'spp_block':
         model = SPPF(in_dim       = in_dim,
         model = SPPF(in_dim       = in_dim,
                      out_dim      = out_dim,
                      out_dim      = out_dim,
-                     expand_ratio = cfg['neck_expand_ratio'],
-                     pooling_size = cfg["spp_pooling_size"],
-                     act_type     = cfg['neck_act'],
-                     norm_type    = cfg['neck_norm']
+                     expand_ratio = cfg.neck_expand_ratio,
+                     pooling_size = cfg.spp_pooling_size,
+                     act_type     = cfg.neck_act,
+                     norm_type    = cfg.neck_norm,
                      )
                      )
         
         
     # ----------------------- FPN Neck -----------------------
     # ----------------------- FPN Neck -----------------------
-    elif cfg['neck'] == 'basic_fpn':
+    elif cfg.neck == 'basic_fpn':
         model = BasicFPN(in_dims = in_dim,
         model = BasicFPN(in_dims = in_dim,
                          out_dim = out_dim,
                          out_dim = out_dim,
-                         p6_feat = cfg['fpn_p6_feat'],
-                         p7_feat = cfg['fpn_p7_feat'],
-                         from_c5 = cfg['fpn_p6_from_c5'], 
+                         p6_feat = cfg.fpn_p6_feat,
+                         p7_feat = cfg.fpn_p7_feat,
+                         from_c5 = cfg.fpn_p6_from_c5, 
                          )
                          )
-    elif cfg['neck'] == 'hybrid_encoder':
-        return HybridEncoder(in_dims     = in_dim,
-                             out_dim     = out_dim,
-                             num_blocks  = cfg['fpn_num_blocks'],
-                             expansion   = cfg['fpn_expansion'],
-                             act_type    = cfg['fpn_act'],
-                             norm_type   = cfg['fpn_norm'],
-                             depthwise   = cfg['fpn_depthwise'],
-                             num_heads   = cfg['en_num_heads'],
-                             num_layers  = cfg['en_num_layers'],
-                             ffn_dim     = cfg['en_ffn_dim'],
-                             dropout     = cfg['en_dropout'],
-                             pe_temperature = cfg['pe_temperature'],
-                             en_act_type    = cfg['en_act'],
-                             en_pre_norm    = cfg['en_pre_norm'],
-                             )
     else:
     else:
-        raise NotImplementedError("Unknown PaFPN: <{}>".format(cfg['fpn']))
+        raise NotImplementedError("Unknown Neck: <{}>".format(cfg.fpn))
         
         
     return model
     return model