diff --git a/train.py b/train.py index 8cb68fc0748e..53dea89694ee 100644 --- a/train.py +++ b/train.py @@ -60,7 +60,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary device, callbacks ): - save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, = \ + save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \ Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \ opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze @@ -124,7 +124,12 @@ 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 + if len(freeze) >= 2: + freeze = freeze.split(',') + freeze = [x.replace(" ", "") for x in freeze] + freeze = [f'model.{x}.' for x in range(int(freeze[0]), int(freeze[1]) + 1)] # layers to freeze + else: + freeze = [f'model.{x}.' for x in range(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 +474,8 @@ 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=str, default='', + help='Specify the layers that you want to freeze in training. (e.g. 10 represents freeze up to 10, or 1,2,4,6 represents freeze layers 1,2,4,6)') 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')