🎨 [Update] logging model place

examples/example_train.py

@@ -13,7 +13,6 @@ from yolo.model.yolo import get_model
 from yolo.tools.log_helper import custom_logger
 from yolo.tools.trainer import Trainer
 from yolo.utils.dataloader import get_dataloader
-from yolo.utils.drawer import draw_model
 from yolo.utils.get_dataset import prepare_dataset
 
 
@@ -24,7 +23,6 @@ def main(cfg: Config):
 
     dataloader = get_dataloader(cfg)
     model = get_model(cfg)
-    draw_model(model=model)
     # TODO: get_device or rank, for DDP mode
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 

yolo/model/yolo.py

@@ -7,6 +7,7 @@ from omegaconf import ListConfig, OmegaConf
 from yolo.config.config import Config, Model, YOLOLayer
 from yolo.tools.layer_helper import get_layer_map
 from yolo.tools.log_helper import log_model
+from yolo.utils.drawer import draw_model
 
 
 class YOLO(nn.Module):
@@ -24,8 +25,6 @@ class YOLO(nn.Module):
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.model: List[YOLOLayer] = nn.ModuleList()
         self.build_model(model_cfg.model)
-        # TODO: Move to other position
-        log_model(self.model)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         self.layer_index = {}
@@ -126,4 +125,6 @@ def get_model(cfg: Config) -> YOLO:
     OmegaConf.set_struct(cfg.model, False)
     model = YOLO(cfg.model, cfg.hyper.data.class_num)
     logger.info("✅ Success load model")
+    log_model(model.model)
+    draw_model(model=model)
     return model