提交 e71fd0ec authored 作者: Glenn Jocher's avatar Glenn Jocher

Model freeze capability (#679)

上级 9ae86836
...@@ -73,6 +73,14 @@ def train(hyp, opt, device, tb_writer=None): ...@@ -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
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论