| from models.classifier import Classifier |
| from models.convnextv2 import convnextv2_tiny, remap_checkpoint_keys, load_state_dict |
| from util.lars import LARS |
| import torch |
| import os |
| from util.convnext_optim import get_parameter_groups, LayerDecayValueAssigner |
|
|
| def build_model(args, device): |
| if args.model == "profound_conv": |
| convnext = convnextv2_tiny(in_chans=3, drop_path_rate=0.1) |
| if args.pretrain is None: |
| raise NotImplementedError(f"No pretrained weight") |
| if not os.path.exists(args.pretrain): |
| raise FileExistsError(f"{args.pretrain} Not exists") |
| ckpt = torch.load(args.pretrain, map_location="cpu") |
| ckpt = remap_checkpoint_keys(ckpt) |
| load_state_dict(convnext, ckpt, weights_only=False) |
| model = Classifier(convnext, args.num_classes) |
| model = model.to(device) |
| if args.train == "freeze": |
| for key, value in model.encoder.named_parameters(): |
| value.requires_grad = False |
| optimizer = LARS(model.head.parameters(), weight_decay=0, lr=args.lr) |
| else: |
| num_layers = sum(convnext.depths) |
| assigner = LayerDecayValueAssigner( |
| list( |
| args.layer_decay ** (num_layers + 1 - i) for i in range(num_layers + 2) |
| ), |
| depths=convnext.depths, |
| layer_decay_type=args.layer_decay_type, |
| ) |
|
|
| skip = {} |
| if hasattr(model.encoder, "no_weight_decay"): |
| skip = model.encoder.no_weight_decay() |
|
|
| backbone_param_groups = get_parameter_groups( |
| model.encoder, |
| args.weight_decay, |
| skip, |
| assigner.get_layer_id, |
| assigner.get_scale, |
| ) |
| decoder_param_groups = [ |
| {"params": model.head.parameters(), "weight_decay": 0.0, "lr": args.lr} |
| ] |
|
|
| optimizer = torch.optim.AdamW( |
| backbone_param_groups + decoder_param_groups, lr=args.lr |
| ) |
|
|
| else: |
| raise NotImplementedError(f"unknown model: {args.model}") |
|
|
| n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) |
|
|
| print("Model = %s" % str(model)) |
| print("number of params (M): %.2f" % (n_parameters / 1.0e6)) |
|
|
| return model, optimizer |
|
|
|
|
| def vit_backbone_parameters( |
| model: torch.nn.Module, weight_decay=1e-5, no_weight_decay_list=(), lr=1e-3 |
| ): |
| no_weight_decay_list = set(no_weight_decay_list) |
| decay = [] |
| no_decay = [] |
|
|
| for name, param in model.named_parameters(): |
| if not param.requires_grad: |
| continue |
|
|
| if param.ndim <= 1 or name.endswith(".bias") or name in no_weight_decay_list: |
| no_decay.append(param) |
| else: |
| decay.append(param) |
|
|
| return [ |
| {"params": no_decay, "weight_decay": 0.0, "lr": lr}, |
| {"params": decay, "weight_decay": weight_decay, "lr": lr}, |
| ] |
|
|
