From 361705d9be532cfff592b8a89db40c8218ed1df2 Mon Sep 17 00:00:00 2001 From: Felix You <35478566+youyuxiansen@users.noreply.github.com> Date: Fri, 17 Dec 2021 22:42:26 +0800 Subject: [PATCH] Multi-layer capable `--freeze` argument (#6019) * support specfiy multiple frozen layers * fix bug * Cleanup Freeze section * Cleanup argument Co-authored-by: Glenn Jocher --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index 8cb68fc..ae19c18 100644 --- a/train.py +++ b/train.py @@ -124,7 +124,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create # Freeze - freeze = [f'model.{x}.' for x in range(freeze)] # layers to freeze + freeze = [f'model.{x}.' for x in (freeze if isinstance(freeze, list) else range(freeze))] # layers to freeze for k, v in model.named_parameters(): v.requires_grad = True # train all layers if any(x in k for x in freeze): @@ -469,7 +469,7 @@ def parse_opt(known=False): parser.add_argument('--linear-lr', action='store_true', help='linear LR') parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon') parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)') - parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24') + parser.add_argument('--freeze', nargs='+', type=int, default=0, help='Freeze layers: backbone=10, first3=0 1 2') parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)') parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')