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

Add `--sync-bn` known issue (#4032) 

* Add `--sync-bn` known issue

* Update train.py
This commit is contained in:
Glenn Jocher 2021-07-17 13:07:19 +02:00 committed by GitHub
parent b7e985e397
commit 951922c735
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
train.py
View File

@ -217,6 +217,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary
# SyncBatchNorm # SyncBatchNorm
if opt.sync_bn and cuda and RANK != -1: if opt.sync_bn and cuda and RANK != -1:
raise Exception('can not train with --sync-bn, known issue https://github.com/ultralytics/yolov5/issues/3998')
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
logger.info('Using SyncBatchNorm()') logger.info('Using SyncBatchNorm()')
@ -232,9 +233,9 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary
# Process 0 # Process 0
if RANK in [-1, 0]: if RANK in [-1, 0]:
valloader = create_dataloader(val_path, imgsz_val, batch_size // WORLD_SIZE * 2, gs, single_cls, valloader = create_dataloader(val_path, imgsz_val, batch_size // WORLD_SIZE * 2, gs, single_cls,
hyp=hyp, cache=opt.cache_images and not noval, rect=True, rank=-1, hyp=hyp, cache=opt.cache_images and not noval, rect=True, rank=-1,
workers=workers, workers=workers,
pad=0.5, prefix=colorstr('val: '))[0] pad=0.5, prefix=colorstr('val: '))[0]
if not resume: if not resume:
labels = np.concatenate(dataset.labels, 0) labels = np.concatenate(dataset.labels, 0)