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

Update train.py forward simplification

pull/860/head
Glenn Jocher 2020-08-25 13:48:03 -07:00
parent 455f7b8f76
commit a21bd0687c
1 changed files with 3 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
train.py
View File

@ -265,18 +265,12 @@ def train(hyp, opt, device, tb_writer=None):
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Autocast
with amp.autocast(enabled=cuda):
# Forward # Forward
pred = model(imgs) with amp.autocast(enabled=cuda):
pred = model(imgs) # forward
# Loss loss, loss_items = compute_loss(pred, targets.to(device), model) # loss scaled by batch_size
loss, loss_items = compute_loss(pred, targets.to(device), model) # scaled by batch_size
if rank != -1: if rank != -1:
loss *= opt.world_size # gradient averaged between devices in DDP mode loss *= opt.world_size # gradient averaged between devices in DDP mode
# if not torch.isfinite(loss):
# logger.info('WARNING: non-finite loss, ending training ', loss_items)
# return results
# Backward # Backward
scaler.scale(loss).backward() scaler.scale(loss).backward()