mirrors
/
yolov5
mirror of https://github.com/ultralytics/yolov5.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Fix missing model.stride in DP and DDP mode (#683)

pull/688/head
NanoCode012 2020-08-10 01:01:36 +07:00 committed by GitHub
parent a0ac5adb7b
commit 3d8ed0a76b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
train.py
View File

@ -124,6 +124,10 @@ def train(hyp, opt, device, tb_writer=None):
del ckpt, state_dict 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 # DP mode
if cuda and rank == -1 and torch.cuda.device_count() > 1: if cuda and rank == -1 and torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model) model = torch.nn.DataParallel(model)
@ -140,10 +144,6 @@ def train(hyp, opt, device, tb_writer=None):
if cuda and rank != -1: if cuda and rank != -1:
model = DDP(model, device_ids=[opt.local_rank], output_device=(opt.local_rank)) 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 # Trainloader
dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True, 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, cache=opt.cache_images, rect=opt.rect, local_rank=rank,