|
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 |
