💬 [Update] logging tools for quite mode

yolo/model/yolo.py

@@ -6,10 +6,8 @@ from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 from torch import nn
 
-from yolo.config.config import Config, ModelConfig, YOLOLayer
+from yolo.config.config import ModelConfig, YOLOLayer
 from yolo.tools.dataset_preparation import prepare_weight
-from yolo.tools.drawer import draw_model
-from yolo.utils.logging_utils import log_model_structure
 from yolo.utils.module_utils import get_layer_map
 
 
@@ -138,6 +136,4 @@ def create_model(model_cfg: ModelConfig, weight_path: Optional[str], class_num:
             model.model.load_state_dict(torch.load(weight_path), strict=False)
             logger.info("✅ Success load model weight")
 
-    log_model_structure(model.model)
-    draw_model(model=model)
     return model

yolo/tools/solver.py

@@ -13,10 +13,10 @@ from torch.utils.data import DataLoader
 from yolo.config.config import Config, TrainConfig, ValidationConfig
 from yolo.model.yolo import YOLO
 from yolo.tools.data_loader import StreamDataLoader, create_dataloader
-from yolo.tools.drawer import draw_bboxes
+from yolo.tools.drawer import draw_bboxes, draw_model
 from yolo.tools.loss_functions import create_loss_function
 from yolo.utils.bounding_box_utils import Vec2Box, bbox_nms, calculate_map
-from yolo.utils.logging_utils import ProgressLogger
+from yolo.utils.logging_utils import ProgressLogger, log_model_structure
 from yolo.utils.model_utils import (
     ExponentialMovingAverage,
     create_optimizer,
@@ -25,7 +25,7 @@ from yolo.utils.model_utils import (
 
 
 class ModelTrainer:
-    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
+    def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device, use_ddp: bool):
         train_cfg: TrainConfig = cfg.task
         self.model = model if not use_ddp else DDP(model, device_ids=[device])
         self.use_ddp = use_ddp
@@ -37,7 +37,13 @@ class ModelTrainer:
         self.progress = progress
         self.num_epochs = cfg.task.epoch
 
-        self.validation_dataloader = create_dataloader(cfg.task.validation.data, cfg.dataset, cfg.task.validation.task)
+        if not progress.quite_mode:
+            log_model_structure(model.model)
+            draw_model(model=model)
+
+        self.validation_dataloader = create_dataloader(
+            cfg.task.validation.data, cfg.dataset, cfg.task.validation.task, use_ddp
+        )
         self.validator = ModelValidator(cfg.task.validation, model, vec2box, progress, device)
 
         if getattr(train_cfg.ema, "enabled", False):

yolo/utils/logging_utils.py

@@ -40,7 +40,9 @@ def custom_logger(quite: bool = False):
 
 class ProgressLogger:
     def __init__(self, cfg: Config, exp_name: str):
-        custom_logger(getattr(cfg, "quite", False))
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        self.quite_mode = local_rank or getattr(cfg, "quite", False)
+        custom_logger(self.quite_mode)
         self.save_path = validate_log_directory(cfg, exp_name=cfg.name)
 
         self.progress = Progress(
@@ -53,7 +55,7 @@ class ProgressLogger:
         if self.use_wandb:
             wandb.errors.term._log = custom_wandb_log
             self.wandb = wandb.init(
-                project="YOLO", resume="allow", mode="online", dir=save_path, id=None, name=exp_name
+                project="YOLO", resume="allow", mode="online", dir=self.save_path, id=None, name=exp_name
             )
 
     def start_train(self, num_epochs: int):