File size: 2,886 Bytes
178d33b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | from typing import Any
import numpy as np
import torch
import torch.nn as nn
from numpy.linalg import norm, pinv
from scipy.special import logsumexp
from sklearn.covariance import EmpiricalCovariance
from tqdm import tqdm
from .base_postprocessor import BasePostprocessor
class VIMPostprocessor(BasePostprocessor):
def __init__(self, config):
super().__init__(config)
self.args = self.config.postprocessor.postprocessor_args
self.args_dict = self.config.postprocessor.postprocessor_sweep
self.dim = self.args.dim
self.setup_flag = False
def setup(self, net: nn.Module, id_loader_dict, ood_loader_dict):
if not self.setup_flag:
net.eval()
with torch.no_grad():
self.w, self.b = net.get_fc()
print('Extracting id training feature')
feature_id_train = []
for batch in tqdm(id_loader_dict['train'],
desc='Setup: ',
position=0,
leave=True):
data = batch['data'].cuda()
data = data.float()
_, feature = net(data, return_feature=True)
feature_id_train.append(feature.cpu().numpy())
feature_id_train = np.concatenate(feature_id_train, axis=0)
logit_id_train = feature_id_train @ self.w.T + self.b
self.u = -np.matmul(pinv(self.w), self.b)
ec = EmpiricalCovariance(assume_centered=True)
ec.fit(feature_id_train - self.u)
eig_vals, eigen_vectors = np.linalg.eig(ec.covariance_)
self.Vim = np.ascontiguousarray(
(eigen_vectors.T[np.argsort(eig_vals * -1)[self.dim:]]).T)
vlogit_id_train = norm(np.matmul(feature_id_train - self.u,
self.NS),
axis=-1)
self.alpha = logit_id_train.max(
axis=-1).mean() / vlogit_id_train.mean()
print(f'self.alpha = {self.alpha:.4f}')
self.setup_flag = True
else:
pass
@torch.no_grad()
def postprocess(self, net: nn.Module, data: Any):
_, feature_ood = net.forward(data, return_feature=True)
feature_ood = feature_ood.cpu()
logit_ood = feature_ood @ self.w.T + self.b
_, pred = torch.max(logit_ood, dim=1)
energy_ood = logsumexp(logit_ood.numpy(), axis=-1)
vlogit_ood = norm(np.matmul(feature_ood.numpy() - self.u, self.NS),
axis=-1) * self.alpha
score_ood = -vlogit_ood + energy_ood
return pred, torch.from_numpy(score_ood)
def set_hyperparam(self, hyperparam: list):
self.dim = hyperparam[0]
def get_hyperparam(self):
return self.dim
|