zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Remove DDP process group timeout (#4422)

This commit is contained in:
Glenn Jocher 2021-08-15 18:32:41 +02:00 committed by GitHub
parent 4e65052f28
commit 19d03a955c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
train.py
View File

@ -493,7 +493,7 @@ def main(opt):
assert not opt.sync_bn, '--sync-bn known training issue, see https://github.com/ultralytics/yolov5/issues/3998' assert not opt.sync_bn, '--sync-bn known training issue, see https://github.com/ultralytics/yolov5/issues/3998'
torch.cuda.set_device(LOCAL_RANK) torch.cuda.set_device(LOCAL_RANK)
device = torch.device('cuda', LOCAL_RANK) device = torch.device('cuda', LOCAL_RANK)
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=60)) dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
# Train # Train
if not opt.evolve: if not opt.evolve:

4
utils/torch_utils.py
View File

@ -35,10 +35,10 @@ def torch_distributed_zero_first(local_rank: int):
Decorator to make all processes in distributed training wait for each local_master to do something. Decorator to make all processes in distributed training wait for each local_master to do something.
""" """
if local_rank not in [-1, 0]: if local_rank not in [-1, 0]:
dist.barrier() dist.barrier(device_ids=[local_rank])
yield yield
if local_rank == 0: if local_rank == 0:
dist.barrier() dist.barrier(device_ids=[0])
def init_torch_seeds(seed=0): def init_torch_seeds(seed=0):