| import os |
| import pickle |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def assign_learning_rate(param_group, new_lr): |
| param_group["lr"] = new_lr |
|
|
|
|
| def _warmup_lr(base_lr, warmup_length, step): |
| return base_lr * (step + 1) / warmup_length |
|
|
|
|
| def cosine_lr(optimizer, base_lrs, warmup_length, steps): |
| if not isinstance(base_lrs, list): |
| base_lrs = [base_lrs for _ in optimizer.param_groups] |
| assert len(base_lrs) == len(optimizer.param_groups) |
|
|
| def _lr_adjuster(step): |
| for param_group, base_lr in zip(optimizer.param_groups, base_lrs): |
| if step < warmup_length: |
| lr = _warmup_lr(base_lr, warmup_length, step) |
| else: |
| e = step - warmup_length |
| es = steps - warmup_length |
| lr = 0.5 * (1 + np.cos(np.pi * e / es)) * base_lr |
| assign_learning_rate(param_group, lr) |
|
|
| return _lr_adjuster |
|
|
|
|
| def accuracy(output, target, topk=(1,)): |
| pred = output.topk(max(topk), 1, True, True)[1].t() |
| correct = pred.eq(target.view(1, -1).expand_as(pred)) |
| return [ |
| float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) |
| for k in topk |
| ] |
|
|
|
|
| def torch_load_old(save_path, device=None): |
| with open(save_path, "rb") as f: |
| classifier = pickle.load(f) |
| if device is not None: |
| classifier = classifier.to(device) |
| return classifier |
|
|
|
|
| def torch_save(model, save_path): |
| if os.path.dirname(save_path) != "": |
| os.makedirs(os.path.dirname(save_path), exist_ok=True) |
| torch.save(model, save_path) |
|
|
|
|
| def torch_load(save_path, device=None): |
| model = torch.load(save_path, map_location="cpu") |
| if device is not None: |
| model = model.to(device) |
| return model |
|
|
|
|
| def get_logits(inputs, classifier): |
| assert callable(classifier) |
| if hasattr(classifier, "to"): |
| classifier = classifier.to(inputs.device) |
| return classifier(inputs) |
|
|
|
|
| def get_probs(inputs, classifier): |
| if hasattr(classifier, "predict_proba"): |
| probs = classifier.predict_proba(inputs.detach().cpu().numpy()) |
| return torch.from_numpy(probs) |
| logits = get_logits(inputs, classifier) |
| return logits.softmax(dim=1) |
|
|
|
|
| class LabelSmoothing(torch.nn.Module): |
| def __init__(self, smoothing=0.0): |
| super(LabelSmoothing, self).__init__() |
| self.confidence = 1.0 - smoothing |
| self.smoothing = smoothing |
|
|
| def forward(self, x, target): |
| logprobs = torch.nn.functional.log_softmax(x, dim=-1) |
|
|
| nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1)) |
| nll_loss = nll_loss.squeeze(1) |
| smooth_loss = -logprobs.mean(dim=-1) |
| loss = self.confidence * nll_loss + self.smoothing * smooth_loss |
| return loss.mean() |
|
|
|
|
| class DotDict(dict): |
| """dot.notation access to dictionary attributes""" |
|
|
| __getattr__ = dict.get |
| __setattr__ = dict.__setitem__ |
| __delattr__ = dict.__delitem__ |
|
|
|
|
| def find_optimal_coef( |
| results, |
| metric="avg_normalized_top1", |
| minimize=False, |
| control_metric=None, |
| control_metric_threshold=0.0, |
| ): |
| best_coef = None |
| if minimize: |
| best_metric = 1 |
| else: |
| best_metric = 0 |
| for scaling_coef in results.keys(): |
| if control_metric is not None: |
| if results[scaling_coef][control_metric] < control_metric_threshold: |
| print(f"Control metric fell below {control_metric_threshold} threshold") |
| continue |
| if minimize: |
| if results[scaling_coef][metric] < best_metric: |
| best_metric = results[scaling_coef][metric] |
| best_coef = scaling_coef |
| else: |
| if results[scaling_coef][metric] > best_metric: |
| best_metric = results[scaling_coef][metric] |
| best_coef = scaling_coef |
| return best_coef |
|
|
|
|
| def nonlinear_advantage(nonlinear_acc, linear_acc, num_classes): |
| return (nonlinear_acc - linear_acc) / (1.0 - 1.0 / num_classes) |
|
|
|
|
| def calculate_linearized_orthogonality_loss(linearized_model): |
| """Compute orthogonality loss ||delta_W^T delta_W - I||_F for a LinearizedModel.""" |
| ortho_loss = 0.0 |
| for p_finetuned, p_pretrained in zip(linearized_model.params, linearized_model.params0): |
| if p_finetuned.requires_grad and p_finetuned.dim() == 2: |
| delta_W = p_finetuned - p_pretrained |
|
|
| rows, cols = delta_W.shape |
| if rows < cols: |
| mat = delta_W @ delta_W.T |
| identity = torch.eye(rows, device=delta_W.device) |
| else: |
| mat = delta_W.T @ delta_W |
| identity = torch.eye(cols, device=delta_W.device) |
|
|
| ortho_loss += torch.norm(mat - identity, p='fro') |
|
|
| return ortho_loss |
|
|
|
|
| def calculate_standard_orthogonality_loss(model, pretrained_state_dict): |
| """Compute orthogonality loss ||delta_W^T delta_W - I||_F for standard/linear-2 finetuning. |
| |
| Args: |
| model: DDP-wrapped ImageClassifier (ddp_model). |
| pretrained_state_dict: snapshot of the pretrained model's inner ViT state_dict. |
| """ |
| ortho_loss = 0.0 |
|
|
| for name, p_finetuned in model.module.image_encoder.model.named_parameters(): |
| if p_finetuned.requires_grad and p_finetuned.dim() == 2: |
| if name in pretrained_state_dict: |
| p_pretrained = pretrained_state_dict[name].to(p_finetuned.device) |
|
|
| delta_W = p_finetuned - p_pretrained |
|
|
| rows, cols = delta_W.shape |
| if rows < cols: |
| mat = delta_W @ delta_W.T |
| identity = torch.eye(rows, device=delta_W.device) |
| else: |
| mat = delta_W.T @ delta_W |
| identity = torch.eye(cols, device=delta_W.device) |
|
|
| ortho_loss += torch.norm(mat - identity, p='fro') |
|
|
| return ortho_loss |
|
|
