updates
This commit is contained in:
parent
0245ff9133
commit
75e8ec323f
13
train.py
13
train.py
|
@ -107,11 +107,14 @@ def train():
|
||||||
chkpt = torch.load(weights, map_location=device)
|
chkpt = torch.load(weights, map_location=device)
|
||||||
|
|
||||||
# load model
|
# load model
|
||||||
# if opt.transfer:
|
try:
|
||||||
chkpt['model'] = {k: v for k, v in chkpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
|
chkpt['model'] = {k: v for k, v in chkpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
|
||||||
model.load_state_dict(chkpt['model'], strict=False)
|
model.load_state_dict(chkpt['model'], strict=False)
|
||||||
# else:
|
# model.load_state_dict(chkpt['model'])
|
||||||
# model.load_state_dict(chkpt['model'])
|
except KeyError as e:
|
||||||
|
s = "%s is not compatible with %s. Specify --weights '' or specify a --cfg compatible with %s. " \
|
||||||
|
"See https://github.com/ultralytics/yolov3/issues/657" % (opt.weights, opt.cfg, opt.weights)
|
||||||
|
raise KeyError(s) from e
|
||||||
|
|
||||||
# load optimizer
|
# load optimizer
|
||||||
if chkpt['optimizer'] is not None:
|
if chkpt['optimizer'] is not None:
|
||||||
|
|
Loading…
Reference in New Issue