--image_weights bug fix (#1524)

This commit is contained in:
Glenn Jocher authored and GitHub committed 2020-11-26 11:49:01 +01:00
1 parent e9a0ae6f19
commit 9728e2b8ae
2 files changed
+10 -8

No files matched your search

+3 -2
View File
@@ -181,8 +181,9 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
# Trainloader
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect,
rank=rank, world_size=opt.world_size, workers=opt.workers)
hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
world_size=opt.world_size, workers=opt.workers,
image_weights=opt.image_weights)
mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class
nb = len(dataloader) # number of batches
assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)