Browse Source

Fix ema attribute error in DDP mode (#775)

* Fix ema error in DDP mode

* Update train.py

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
5.0
NanoCode012 GitHub 4 years ago
parent
commit
fb4fc8cd02
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      train.py

+ 1
- 1
train.py View File

@@ -164,11 +164,11 @@ def train(hyp, opt, device, tb_writer=None):
world_size=opt.world_size, workers=opt.workers)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
ema.updates = start_epoch * nb // accumulate # set EMA updates
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)

# Testloader
if rank in [-1, 0]:
ema.updates = start_epoch * nb // accumulate # set EMA updates
testloader = create_dataloader(test_path, imgsz_test, total_batch_size, gs, opt,
hyp=hyp, augment=False, cache=opt.cache_images, rect=True, rank=-1,
world_size=opt.world_size, workers=opt.workers)[0] # only runs on process 0

Loading…
Cancel
Save