diff --git a/train.py b/train.py index 7953308d4b32..03706eeb1e2a 100644 --- a/train.py +++ b/train.py @@ -221,8 +221,9 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Trainloader train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls, - hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=LOCAL_RANK, - workers=workers, image_weights=opt.image_weights, quad=opt.quad, + hyp=hyp, augment=True, cache=None if opt.cache == 'val' else opt.cache, + rect=opt.rect, rank=LOCAL_RANK, workers=workers, + image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '), shuffle=True) mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class nb = len(train_loader) # number of batches @@ -231,8 +232,8 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Process 0 if RANK in [-1, 0]: val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls, - hyp=hyp, cache=None if noval else opt.cache, rect=True, rank=-1, - workers=workers * 2, pad=0.5, + hyp=hyp, cache=None if noval else opt.cache, + rect=True, rank=-1, workers=workers * 2, pad=0.5, prefix=colorstr('val: '))[0] if not resume: