Skip to content

Commit

Permalink
Fix missing model.stride in DP and DDP mode (ultralytics#683)
Browse files Browse the repository at this point in the history
  • Loading branch information
NanoCode012 committed Aug 9, 2020
1 parent c96b169 commit bcedb91
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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,
Expand Down

0 comments on commit bcedb91

Please sign in to comment.