diff --git a/train.py b/train.py index 39b996c34a38..8c07ec2a744f 100644 --- a/train.py +++ b/train.py @@ -123,6 +123,10 @@ def train(hyp, opt, device, tb_writer=None): epochs += ckpt['epoch'] # finetune additional epochs del ckpt, state_dict + + # Image sizes + gs = int(max(model.stride)) # grid size (max stride) + imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples # DP mode if cuda and rank == -1 and torch.cuda.device_count() > 1: @@ -140,10 +144,6 @@ def train(hyp, opt, device, tb_writer=None): if cuda and rank != -1: model = DDP(model, device_ids=[opt.local_rank], output_device=(opt.local_rank)) - # Image sizes - gs = int(max(model.stride)) # grid size (max stride) - imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size] # verify imgsz are gs-multiples - # Trainloader dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, local_rank=rank,