🚀 [New] DDP mode for training model

yolo/lazy.py

@@ -2,7 +2,6 @@ import sys
 from pathlib import Path
 
 import hydra
-import torch
 
 project_root = Path(__file__).resolve().parent.parent
 sys.path.append(str(project_root))
@@ -14,22 +13,24 @@ from yolo.tools.solver import ModelTester, ModelTrainer
 from yolo.utils.bounding_box_utils import Vec2Box
 from yolo.utils.deploy_utils import FastModelLoader
 from yolo.utils.logging_utils import ProgressLogger
-from yolo.utils.model_utils import send_to_device
+from yolo.utils.model_utils import get_device
 
 
 @hydra.main(config_path="config", config_name="config", version_base=None)
 def main(cfg: Config):
     progress = ProgressLogger(cfg, exp_name=cfg.name)
-    dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task)
+    device, use_ddp = get_device(cfg.device)
+    dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
     if getattr(cfg.task, "fast_inference", False):
         model = FastModelLoader(cfg).load_model()
     else:
         model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight)
-    device, model = send_to_device(model, cfg.device)
+        model = model.to(device)
+
     vec2box = Vec2Box(model, cfg.image_size, device)
 
     if cfg.task.task == "train":
-        trainer = ModelTrainer(cfg, model, vec2box, progress, device)
+        trainer = ModelTrainer(cfg, model, vec2box, progress, device, use_ddp)
         trainer.solve(dataloader)
 
     if cfg.task.task == "inference":

yolo/tools/data_loader.py

@@ -12,7 +12,7 @@ from PIL import Image
 from rich.progress import track
 from torch import Tensor
 from torch.utils.data import DataLoader, Dataset
-from torchvision.transforms import functional as TF
+from torch.utils.data.distributed import DistributedSampler
 
 from yolo.config.config import DataConfig, DatasetConfig
 from yolo.tools.data_augmentation import (
@@ -157,14 +157,16 @@ class YoloDataset(Dataset):
 
 
 class YoloDataLoader(DataLoader):
-    def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train"):
+    def __init__(self, data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):
         """Initializes the YoloDataLoader with hydra-config files."""
         dataset = YoloDataset(data_cfg, dataset_cfg, task)
+        sampler = DistributedSampler(dataset, shuffle=data_cfg.shuffle) if use_ddp else None
         self.image_size = data_cfg.image_size[0]
         super().__init__(
             dataset,
             batch_size=data_cfg.batch_size,
-            shuffle=data_cfg.shuffle,
+            sampler=sampler,
+            shuffle=data_cfg.shuffle and not use_ddp,
             num_workers=data_cfg.cpu_num,
             pin_memory=data_cfg.pin_memory,
             collate_fn=self.collate_fn,
@@ -198,14 +200,14 @@ class YoloDataLoader(DataLoader):
         return batch_images, batch_targets
 
 
-def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train"):
+def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):
     if task == "inference":
         return StreamDataLoader(data_cfg)
 
     if dataset_cfg.auto_download:
         prepare_dataset(dataset_cfg, task)
 
-    return YoloDataLoader(data_cfg, dataset_cfg, task)
+    return YoloDataLoader(data_cfg, dataset_cfg, task, use_ddp)
 
 
 class StreamDataLoader:

yolo/tools/solver.py

@@ -7,6 +7,8 @@ from torch import Tensor
 
 # TODO: We may can't use CUDA?
 from torch.cuda.amp import GradScaler, autocast
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader
 
 from yolo.config.config import Config, TrainConfig, ValidationConfig
 from yolo.model.yolo import YOLO
@@ -25,7 +27,8 @@ from yolo.utils.model_utils import (
 class ModelTrainer:
     def __init__(self, cfg: Config, model: YOLO, vec2box: Vec2Box, progress: ProgressLogger, device):
         train_cfg: TrainConfig = cfg.task
-        self.model = model
+        self.model = model if not use_ddp else DDP(model, device_ids=[device])
+        self.use_ddp = use_ddp
         self.vec2box = vec2box
         self.device = device
         self.optimizer = create_optimizer(model, train_cfg.optimizer)
@@ -86,13 +89,15 @@ class ModelTrainer:
             self.ema.restore()
         torch.save(checkpoint, filename)
 
-    def solve(self, dataloader):
+    def solve(self, dataloader: DataLoader):
         logger.info("🚄 Start Training!")
         num_epochs = self.num_epochs
 
         with self.progress.progress:
             self.progress.start_train(num_epochs)
             for epoch in range(num_epochs):
+                if self.use_ddp:
+                    dataloader.sampler.set_epoch(epoch)
 
                 self.progress.start_one_epoch(len(dataloader), self.optimizer, epoch)
                 epoch_loss = self.train_one_epoch(dataloader)

yolo/utils/model_utils.py

@@ -1,7 +1,9 @@
-from typing import Any, Dict, List, Type, Union
+import os
+from typing import List, Type, Union
 
 import torch
 import torch.distributed as dist
+from loguru import logger
 from omegaconf import ListConfig
 from torch import nn
 from torch.nn.parallel import DistributedDataParallel as DDP
@@ -73,29 +75,21 @@ def create_scheduler(optimizer: Optimizer, schedule_cfg: SchedulerConfig) -> _LR
     return schedule
 
 
-def get_device():
-    if torch.cuda.is_available():
-        return torch.device("cuda")
-    elif torch.backends.mps.is_available():
-        return torch.device("mps")
-    else:
-        return torch.device("cpu")
-
-
-def send_to_device(model: nn.Module, device: Union[str, int, List[int]]):
-    if not isinstance(device, (List, ListConfig)):
-        device = torch.device(device)
-        print("runing man")
-        return device, model.to(device)
-
-    device = torch.device("cuda")
-    world_size = dist.get_world_size()
-    print("runing man")
-    dist.init_process_group(
-        backend="gloo" if torch.cuda.is_available() else "gloo", rank=dist.get_rank(), world_size=world_size
-    )
-    print(f"Initialized process group; rank: {dist.get_rank()}, size: {world_size}")
-
-    model = model.cuda(device)
-    model = DDP(model, device_ids=[device])
-    return device, model.to(device)
+def initialize_distributed() -> None:
+    rank = int(os.getenv("RANK", "0"))
+    local_rank = int(os.getenv("LOCAL_RANK", "0"))
+    world_size = int(os.getenv("WORLD_SIZE", "1"))
+
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    logger.info(f"Initialized process group; rank: {rank}, size: {world_size}")
+    return local_rank
+
+
+def get_device(device_spec: Union[str, int, List[int]]) -> torch.device:
+    ddp_flag = False
+    if isinstance(device_spec, (list, ListConfig)):
+        ddp_flag = True
+        device_spec = initialize_distributed()
+    device = torch.device(device_spec)
+    return device, ddp_flag