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

train.py GPU memory fix (#3590) 

* train.py GPU memory fix

* ema

* cuda

* cuda

* zeros input

* to device

* batch index 0
pull/3593/head
Glenn Jocher 2021-06-11 20:24:03 +02:00 committed by GitHub
parent e8c5237403
commit 4984cf54be
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
train.py
View File

@ -335,7 +335,7 @@ def train(hyp, opt, device, tb_writer=None):
if tb_writer and ni == 0:
with warnings.catch_warnings():
warnings.simplefilter('ignore') # suppress jit trace warning
tb_writer.add_graph(torch.jit.trace(de_parallel(model), imgs, strict=False), []) # graph
tb_writer.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), [])
elif plots and ni == 10 and wandb_logger.wandb:
wandb_logger.log({'Mosaics': [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
save_dir.glob('train*.jpg') if x.exists()]})