ml_lib/metrics/attention_rollout.py
2021-02-15 11:39:54 +01:00

48 lines
1.4 KiB
Python

import numpy as np
from einops import reduce
import torch
from sklearn.ensemble import IsolationForest
from sklearn.metrics import recall_score, roc_auc_score, average_precision_score
from ml_lib.metrics._base_score import _BaseScores
class AttentionRollout(_BaseScores):
def __init__(self, *args):
super(AttentionRollout, self).__init__(*args)
pass
def __call__(self, outputs):
summary_dict = dict()
#######################################################################################
# Additional Score - Histogram Distances - Image Plotting
#######################################################################################
#
# INIT
attn_weights = [output['attn_weights'].cpu().numpy() for output in outputs]
attn_reduce_heads = [reduce(x, '') for x in attn_weights]
if self.model.params.use_residual:
residual_att = np.eye(att_mat.shape[1])[None, ...]
aug_att_mat = att_mat + residual_att
aug_att_mat = aug_att_mat / aug_att_mat.sum(axis=-1)[..., None]
else:
aug_att_mat = att_mat
joint_attentions = np.zeros(aug_att_mat.shape)
layers = joint_attentions.shape[0]
joint_attentions[0] = aug_att_mat[0]
for i in np.arange(1, layers):
joint_attentions[i] = aug_att_mat[i].dot(joint_attentions[i - 1])