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

Improve trainer DDP device handling (#15383)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
alanZee 2024-08-13 04:53:48 +08:00 • committed by GitHub
parent 592f421663
commit 03e0b1033c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/engine/trainer.py
View file

@ -174,9 +174,11 @@ class BaseTrainer:
world_size = len(self.args.device.split(","))
elif isinstance(self.args.device, (tuple, list)): # i.e. device=[0, 1, 2, 3] (multi-GPU from CLI is list)
world_size = len(self.args.device)
elif self.args.device in {"cpu", "mps"}: # i.e. device='cpu' or 'mps'
world_size = 0
elif torch.cuda.is_available(): # i.e. device=None or device='' or device=number
world_size = 1 # default to device 0
else: # i.e. device='cpu' or 'mps'
else: # i.e. device=None or device=''
world_size = 0
# Run subprocess if DDP training, else train normally