From c8e51812a527eef8ad34fd3530b4942ad156b71e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 4 Sep 2020 13:13:10 -0700 Subject: [PATCH] hyp evolution force-autoanchor fix --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index 84728e850..94f3673c4 100644 --- a/train.py +++ b/train.py @@ -68,10 +68,10 @@ def train(hyp, opt, device, tb_writer=None): with torch_distributed_zero_first(rank): attempt_download(weights) # download if not found locally ckpt = torch.load(weights, map_location=device) # load checkpoint - if 'anchors' in hyp and hyp['anchors']: + if hyp.get('anchors'): ckpt['model'].yaml['anchors'] = round(hyp['anchors']) # force autoanchor model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc).to(device) # create - exclude = ['anchor'] if opt.cfg else [] # exclude keys + exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [] # exclude keys state_dict = ckpt['model'].float().state_dict() # to FP32 state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(state_dict, strict=False) # load