Model freeze capability (#679)

pull/1/head
Glenn Jocher 5 years ago
parent 9ae868364a
commit e71fd0ec0b

@ -73,6 +73,14 @@ def train(hyp, opt, device, tb_writer=None):
else: else:
model = Model(opt.cfg, ch=3, nc=nc).to(device) # create model = Model(opt.cfg, ch=3, nc=nc).to(device) # create
# Freeze
freeze = ['', ] # parameter names to freeze (full or partial)
if any(freeze):
for k, v in model.named_parameters():
if any(x in k for x in freeze):
print('freezing %s' % k)
v.requires_grad = False
# Optimizer # Optimizer
nbs = 64 # nominal batch size nbs = 64 # nominal batch size
accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing

Loading…
Cancel
Save