| import numpy as np |
|
|
| import torch |
| import torch.nn as nn |
| from torchvision import models |
|
|
| from scipy.optimize import root_scalar |
| from scipy.special import betainc |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| def build_backbone(path, name='resnet50'): |
| """ Builds a pretrained ResNet-50 backbone. """ |
| model = getattr(models, name)(pretrained=False) |
| model.head = nn.Identity() |
| model.fc = nn.Identity() |
| checkpoint = torch.load(path, map_location=device) |
| state_dict = checkpoint |
| for ckpt_key in ['state_dict', 'model_state_dict', 'teacher']: |
| if ckpt_key in checkpoint: |
| state_dict = checkpoint[ckpt_key] |
| state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} |
| state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()} |
| msg = model.load_state_dict(state_dict, strict=False) |
| return model |
|
|
| def get_linear_layer(weight, bias): |
| """ Creates a layer that performs feature whitening or centering """ |
| dim_out, dim_in = weight.shape |
| layer = nn.Linear(dim_in, dim_out) |
| layer.weight = nn.Parameter(weight) |
| layer.bias = nn.Parameter(bias) |
| return layer |
|
|
| def load_normalization_layer(path): |
| """ |
| Loads the normalization layer from a checkpoint and returns the layer. |
| """ |
| checkpoint = torch.load(path, map_location=device) |
| if 'whitening' in path or 'out' in path: |
| D = checkpoint['weight'].shape[1] |
| weight = torch.nn.Parameter(D*checkpoint['weight']) |
| bias = torch.nn.Parameter(D*checkpoint['bias']) |
| else: |
| weight = checkpoint['weight'] |
| bias = checkpoint['bias'] |
| return get_linear_layer(weight, bias).to(device, non_blocking=True) |
|
|
| class NormLayerWrapper(nn.Module): |
| """ |
| Wraps backbone model and normalization layer |
| """ |
| def __init__(self, backbone, head): |
| super(NormLayerWrapper, self).__init__() |
| backbone.eval(), head.eval() |
| self.backbone = backbone |
| self.head = head |
|
|
| def forward(self, x): |
| output = self.backbone(x) |
| return self.head(output) |
|
|
| def cosine_pvalue(c, d, k=1): |
| """ |
| Returns the probability that the absolute value of the projection |
| between random unit vectors is higher than c |
| Args: |
| c: cosine value |
| d: dimension of the features |
| k: number of dimensions of the projection |
| """ |
| assert k>0 |
| a = (d - k) / 2.0 |
| b = k / 2.0 |
| if c < 0: |
| return 1.0 |
| return betainc(a, b, 1 - c ** 2) |
|
|
| def pvalue_angle(dim, k=1, angle=None, proba=None): |
| def f(a): |
| return cosine_pvalue(np.cos(a), dim, k) - proba |
| a = root_scalar(f, x0=0.49*np.pi, bracket=[0, np.pi/2]) |
| |
| return a.root |
