Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

update train.py ckpt loading

Browse Source
This commit is contained in:
Glenn Jocher committed 2020-07-18 14:34:12 -07:00
1 parent 7b0c9be837
commit 65857adf52
1 file changed
+1 -2
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+1 -2
View File
@@ -119,8 +119,7 @@ def train(hyp):
# load model # load model
try: try:
ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items() ckpt['model'] = {k: v for k, v in ckpt['model'].float().state_dict().items() if k in model.state_dict()}
if model.state_dict()[k].shape == v.shape} # to FP32, filter
model.load_state_dict(ckpt['model'], strict=False) model.load_state_dict(ckpt['model'], strict=False)
except KeyError as e: except KeyError as e:
s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \ s = "%s is not compatible with %s. This may be due to model differences or %s may be out of date. " \