Spaces:
Running
Running
| import math | |
| def adjust_learning_rate(optimizer, | |
| base_lr, | |
| p, | |
| itr, | |
| max_itr, | |
| restart=1, | |
| warm_up_steps=1000, | |
| is_cosine_decay=False, | |
| min_lr=1e-5, | |
| encoder_lr_ratio=1.0, | |
| freeze_params=[]): | |
| if restart > 1: | |
| each_max_itr = int(math.ceil(float(max_itr) / restart)) | |
| itr = itr % each_max_itr | |
| warm_up_steps /= restart | |
| max_itr = each_max_itr | |
| if itr < warm_up_steps: | |
| now_lr = min_lr + (base_lr - min_lr) * itr / warm_up_steps | |
| else: | |
| itr = itr - warm_up_steps | |
| max_itr = max_itr - warm_up_steps | |
| if is_cosine_decay: | |
| now_lr = min_lr + (base_lr - min_lr) * (math.cos(math.pi * itr / | |
| (max_itr + 1)) + | |
| 1.) * 0.5 | |
| else: | |
| now_lr = min_lr + (base_lr - min_lr) * (1 - itr / (max_itr + 1))**p | |
| for param_group in optimizer.param_groups: | |
| if encoder_lr_ratio != 1.0 and "encoder." in param_group["name"]: | |
| param_group['lr'] = (now_lr - min_lr) * encoder_lr_ratio + min_lr | |
| else: | |
| param_group['lr'] = now_lr | |
| for freeze_param in freeze_params: | |
| if freeze_param in param_group["name"]: | |
| param_group['lr'] = 0 | |
| param_group['weight_decay'] = 0 | |
| break | |
| return now_lr | |
| def get_trainable_params(model, | |
| base_lr, | |
| weight_decay, | |
| use_frozen_bn=False, | |
| exclusive_wd_dict={}, | |
| no_wd_keys=[]): | |
| params = [] | |
| memo = set() | |
| total_param = 0 | |
| for key, value in model.named_parameters(): | |
| if value in memo: | |
| continue | |
| total_param += value.numel() | |
| if not value.requires_grad: | |
| continue | |
| memo.add(value) | |
| wd = weight_decay | |
| for exclusive_key in exclusive_wd_dict.keys(): | |
| if exclusive_key in key: | |
| wd = exclusive_wd_dict[exclusive_key] | |
| break | |
| if len(value.shape) == 1: # normalization layers | |
| if 'bias' in key: # bias requires no weight decay | |
| wd = 0. | |
| elif not use_frozen_bn: # if not use frozen BN, apply zero weight decay | |
| wd = 0. | |
| elif 'encoder.' not in key: # if use frozen BN, apply weight decay to all frozen BNs in the encoder | |
| wd = 0. | |
| else: | |
| for no_wd_key in no_wd_keys: | |
| if no_wd_key in key: | |
| wd = 0. | |
| break | |
| params += [{ | |
| "params": [value], | |
| "lr": base_lr, | |
| "weight_decay": wd, | |
| "name": key | |
| }] | |
| print('Total Param: {:.2f}M'.format(total_param / 1e6)) | |
| return params | |
| def freeze_params(module): | |
| for p in module.parameters(): | |
| p.requires_grad = False | |
| def calculate_params(state_dict): | |
| memo = set() | |
| total_param = 0 | |
| for key, value in state_dict.items(): | |
| if value in memo: | |
| continue | |
| memo.add(value) | |
| total_param += value.numel() | |
| print('Total Param: {:.2f}M'.format(total_param / 1e6)) | |