import libmr import numpy as np import scipy.spatial.distance as spd import torch import torch.nn as nn from tqdm import tqdm from .base_postprocessor import BasePostprocessor from .info import num_classes_dict class OpenMax(BasePostprocessor): def __init__(self, config): super(OpenMax, self).__init__(config) self.nc = num_classes_dict[config.dataset.name] self.weibull_alpha = 3 self.weibull_threshold = 0.9 self.weibull_tail = 20 self.setup_flag = False def setup(self, net: nn.Module, id_loader_dict, ood_loder_dict): if not self.setup_flag: # Fit the weibull distribution from training data. print('Fittting Weibull distribution...') _, mavs, dists = compute_train_score_and_mavs_and_dists( self.nc, id_loader_dict['train'], device='cuda', net=net) categories = list(range(0, self.nc)) self.weibull_model = fit_weibull(mavs, dists, categories, self.weibull_tail, 'euclidean') self.setup_flag = True else: pass @torch.no_grad() def postprocess(self, net: nn.Module, data): net.eval() scores = net(data).cpu().numpy() scores = np.array(scores)[:, np.newaxis, :] categories = list(range(0, self.nc)) pred_openmax = [] score_openmax = [] for score in scores: so, _ = openmax(self.weibull_model, categories, score, 0.5, self.weibull_alpha, 'euclidean') # openmax_prob, softmax_prob pred_openmax.append( np.argmax(so) if np.max(so) >= self.weibull_threshold else ( self.nc - 1)) score_openmax.append(so) pred = torch.tensor(pred_openmax) conf = -1 * torch.from_numpy(np.array(score_openmax))[:, -1] return pred, conf def compute_channel_distances(mavs, features, eu_weight=0.5): """ Input: mavs (channel, C) features: (N, channel, C) Output: channel_distances: dict of distance distribution from MAV for each channel. """ eucos_dists, eu_dists, cos_dists = [], [], [] for channel, mcv in enumerate(mavs): # Compute channel specific distances eu_dists.append( [spd.euclidean(mcv, feat[channel]) for feat in features]) cos_dists.append([spd.cosine(mcv, feat[channel]) for feat in features]) eucos_dists.append([ spd.euclidean(mcv, feat[channel]) * eu_weight + spd.cosine(mcv, feat[channel]) for feat in features ]) return { 'eucos': np.array(eucos_dists), 'cosine': np.array(cos_dists), 'euclidean': np.array(eu_dists) } def compute_train_score_and_mavs_and_dists(train_class_num, trainloader, device, net): scores = [[] for _ in range(train_class_num)] train_dataiter = iter(trainloader) with torch.no_grad(): for train_step in tqdm(range(1, len(train_dataiter) + 1), desc='Progress: ', position=0, leave=True): batch = next(train_dataiter) data = batch['data'].cuda() target = batch['label'].cuda() # this must cause error for cifar outputs = net(data) for score, t in zip(outputs, target): if torch.argmax(score) == t: scores[t].append(score.unsqueeze(dim=0).unsqueeze(dim=0)) scores = [torch.cat(x).cpu().numpy() for x in scores] # (N_c, 1, C) * C mavs = np.array([np.mean(x, axis=0) for x in scores]) # (C, 1, C) dists = [ compute_channel_distances(mcv, score) for mcv, score in zip(mavs, scores) ] return scores, mavs, dists def fit_weibull(means, dists, categories, tailsize=20, distance_type='eucos'): """ Input: means (C, channel, C) dists (N_c, channel, C) * C Output: weibull_model : Perform EVT based analysis using tails of distances and save weibull model parameters for re-adjusting softmax scores """ weibull_model = {} for mean, dist, category_name in zip(means, dists, categories): weibull_model[category_name] = {} weibull_model[category_name]['distances_{}'.format( distance_type)] = dist[distance_type] weibull_model[category_name]['mean_vec'] = mean weibull_model[category_name]['weibull_model'] = [] for channel in range(mean.shape[0]): mr = libmr.MR() tailtofit = np.sort(dist[distance_type][channel, :])[-tailsize:] mr.fit_high(tailtofit, len(tailtofit)) weibull_model[category_name]['weibull_model'].append(mr) return weibull_model def compute_openmax_prob(scores, scores_u): prob_scores, prob_unknowns = [], [] for s, su in zip(scores, scores_u): channel_scores = np.exp(s) channel_unknown = np.exp(np.sum(su)) total_denom = np.sum(channel_scores) + channel_unknown prob_scores.append(channel_scores / total_denom) prob_unknowns.append(channel_unknown / total_denom) # Take channel mean scores = np.mean(prob_scores, axis=0) unknowns = np.mean(prob_unknowns, axis=0) modified_scores = scores.tolist() + [unknowns] return modified_scores def query_weibull(category_name, weibull_model, distance_type='eucos'): return [ weibull_model[category_name]['mean_vec'], weibull_model[category_name]['distances_{}'.format(distance_type)], weibull_model[category_name]['weibull_model'] ] def calc_distance(query_score, mcv, eu_weight, distance_type='eucos'): if distance_type == 'eucos': query_distance = spd.euclidean(mcv, query_score) * eu_weight + \ spd.cosine(mcv, query_score) elif distance_type == 'euclidean': query_distance = spd.euclidean(mcv, query_score) elif distance_type == 'cosine': query_distance = spd.cosine(mcv, query_score) else: print('distance type not known: enter either of eucos, \ euclidean or cosine') return query_distance def softmax(x): e_x = np.exp(x - np.max(x)) return e_x / e_x.sum() def openmax(weibull_model, categories, input_score, eu_weight, alpha=10, distance_type='eucos'): """Re-calibrate scores via OpenMax layer Output: openmax probability and softmax probability """ nb_classes = len(categories) ranked_list = input_score.argsort().ravel()[::-1][:alpha] alpha_weights = [((alpha + 1) - i) / float(alpha) for i in range(1, alpha + 1)] omega = np.zeros(nb_classes) omega[ranked_list] = alpha_weights scores, scores_u = [], [] for channel, input_score_channel in enumerate(input_score): score_channel, score_channel_u = [], [] for c, category_name in enumerate(categories): mav, dist, model = query_weibull(category_name, weibull_model, distance_type) channel_dist = calc_distance(input_score_channel, mav[channel], eu_weight, distance_type) wscore = model[channel].w_score(channel_dist) modified_score = input_score_channel[c] * (1 - wscore * omega[c]) score_channel.append(modified_score) score_channel_u.append(input_score_channel[c] - modified_score) scores.append(score_channel) scores_u.append(score_channel_u) scores = np.asarray(scores) scores_u = np.asarray(scores_u) openmax_prob = np.array(compute_openmax_prob(scores, scores_u)) softmax_prob = softmax(np.array(input_score.ravel())) return openmax_prob, softmax_prob