Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2019-12-30 16:13:06 -08:00
parent 9dd1316a70
commit 6290f9fdb7
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
train.py
View File

@ -137,7 +137,7 @@ def train():
cutoff = load_darknet_weights(model, weights) cutoff = load_darknet_weights(model, weights)
if opt.transfer or opt.prebias: # transfer learning edge (yolo) layers if opt.transfer or opt.prebias: # transfer learning edge (yolo) layers
nf = int(model.module_defs[model.yolo_layers[0] - 1]['filters']) # yolo layer size (i.e. 255) nf = [int(model.module_defs[x - 1]['filters']) for x in model.yolo_layers] # yolo layer size (i.e. 255)
if opt.prebias: if opt.prebias:
for p in optimizer.param_groups: for p in optimizer.param_groups:
@ -147,9 +147,9 @@ def train():
p['momentum'] = 0.9 p['momentum'] = 0.9
for p in model.parameters(): for p in model.parameters():
if opt.prebias and p.numel() == nf: # train (yolo biases) if opt.prebias and p.numel() in nf: # train (yolo biases)
p.requires_grad = True p.requires_grad = True
elif opt.transfer and p.shape[0] == nf: # train (yolo biases+weights) elif opt.transfer and p.shape[0] in nf: # train (yolo biases+weights)
p.requires_grad = True p.requires_grad = True
else: # freeze layer else: # freeze layer
p.requires_grad = False p.requires_grad = False