HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.2.71 Multinode DDP training (#14879)

Browse source 
Co-authored-by: Haris Rehman <haris.rehman.cowlar@gmail.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Haris Rehman 2024-08-01 20:31:03 +05:00 • committed by GitHub
parent 16fc325308
commit 9c5d1a2451
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 9 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

5
ultralytics/engine/trainer.py
View file

@ -26,6 +26,7 @@ from ultralytics.data.utils import check_cls_dataset, check_det_dataset
from ultralytics.nn.tasks import attempt_load_one_weight, attempt_load_weights
from ultralytics.utils import (
DEFAULT_CFG,
LOCAL_RANK,
LOGGER,
RANK,
TQDM,
@ -129,7 +130,7 @@ class BaseTrainer:
# Model and Dataset
self.model = check_model_file_from_stem(self.args.model) # add suffix, i.e. yolov8n -> yolov8n.pt
with torch_distributed_zero_first(RANK): # avoid auto-downloading dataset multiple times
with torch_distributed_zero_first(LOCAL_RANK): # avoid auto-downloading dataset multiple times
self.trainset, self.testset = self.get_dataset()
self.ema = None
@ -285,7 +286,7 @@ class BaseTrainer:
# Dataloaders
batch_size = self.batch_size // max(world_size, 1)
self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=RANK, mode="train")
self.train_loader = self.get_dataloader(self.trainset, batch_size=batch_size, rank=LOCAL_RANK, mode="train")
if RANK in {-1, 0}:
# Note: When training DOTA dataset, double batch size could get OOM on images with >2000 objects.
self.test_loader = self.get_dataloader(