From bcedb91f903d6e1d3eb899e36f357cefa8ffea02 Mon Sep 17 00:00:00 2001 From: NanoCode012 Date: Mon, 10 Aug 2020 01:01:36 +0700 Subject: [PATCH] Fix missing model.stride in DP and DDP mode (#683) --- train.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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,