Commit ed3aae19 authored by Alexander Froch's avatar Alexander Froch
Browse files

Adding dips-attention evaluation

parent 682766e0
from umami.configuration import global_config, logger # isort:skip
import argparse
import os
import pickle
import h5py
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.layers import Lambda
from tensorflow.keras.models import Model, load_model
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import CustomObjectScope
import umami.preprocessing_tools as upt
import umami.train_tools as utt
from umami.evaluation_tools.PlottingFunctions import (
GetScore,
GetScoreC,
discriminant_output_shape,
get_gradients,
getDiscriminant,
)
from umami.evaluation_tools.PlottingFunctions import GetScore, GetScoreC
from umami.preprocessing_tools import Configuration
from umami.tf_tools import (
AttentionPooling,
ConditionalAttention,
ConditionalDeepSet,
DenseNet,
)
from umami.train_tools import Sum
# from plottingFunctions import sigBkgEff
tf.compat.v1.disable_eager_execution()
......@@ -408,20 +407,31 @@ def EvaluateModelDips(
logger.info(f"Evaluating {model_file}")
# Get the testfile with the needed configs
X_test_trk, Y_test = utt.GetTestSampleTrks(
X_test, X_test_trk, Y_test = utt.GetTestFile(
test_file,
train_config.var_dict,
preprocess_config,
nJets=nJets,
exclude=None,
)
X_test = X_test[["absEta_btagJes", "pt_btagJes"]]
# Load the model for evaluation. Note: The Sum is needed here!
with CustomObjectScope({"Sum": utt.Sum}):
model = load_model(model_file)
model = load_model(
model_file,
{
"Sum": Sum,
"ConditionalAttention": ConditionalAttention,
"ConditionalDeepSet": ConditionalDeepSet,
"DenseNet": DenseNet,
"AttentionPooling": AttentionPooling,
},
)
# Get predictions from trained model
pred_dips = model.predict(
X_test_trk,
[X_test_trk, X_test],
batch_size=train_config.NN_structure["batch_size"],
verbose=0,
)
......@@ -600,77 +610,77 @@ def EvaluateModelDips(
f.attrs["N_test"] = len(df)
f.close()
logger.info("Calculate gradients for inputs")
# Cut off last layer of the model for saliency maps
cutted_model = model.layers[-1].output
# Define the last node for the discriminant output
disc = Lambda(getDiscriminant, output_shape=discriminant_output_shape)(
cutted_model
)
# Define the computation graph for the model
model = Model(model.inputs, disc)
# Define boolean mask to filter placeholder tracks
boolMask = (np.sum(X_test_trk, axis=-1) != 0).astype(bool)
# Define the number of true tracks per jet as a mask
nTrks = np.sum(boolMask, axis=-1)
# Get score for the dips prediction
Db = GetScore(
pb=pred_dips[:, b_index],
pc=pred_dips[:, c_index],
pu=pred_dips[:, u_index],
)
# Init small dict
map_dict = {}
# Iterate over different beff, jet flavours and passed options
for target_beff in [60, 70, 77, 85]:
for jet_flavour in [0, 1, 2]:
for PassBool in [True, False]:
# Get the Db value for a specific flavour
Db_flavour = Db[Y_test[:, jet_flavour].astype(bool)]
# Get the cutvalue for the specific WP
cutvalue = np.percentile(Db_flavour, (100 - target_beff))
# Set PassBool masking
if PassBool is True:
mask = Y_test[:, jet_flavour].astype(bool)
mask = mask & (nTrks == 8)
mask = mask & (Db > cutvalue)
elif PassBool is False:
mask = Y_test[:, jet_flavour].astype(bool)
mask = mask & (nTrks == 8)
mask = mask & (Db < cutvalue)
# Get gradient map
gradient_map = get_gradients(model, X_test_trk[mask], 15000)
# Turn gradient map for plotting
gradient_map = np.swapaxes(gradient_map, 1, 2)
# Mean over the jets
gradient_map = np.mean(gradient_map, axis=0)
map_dict.update(
{f"{target_beff}_{jet_flavour}_{PassBool}": gradient_map}
)
# Create results dir and .h5 ile
os.system(f"mkdir -p {train_config.model_name}/results")
with open(
f"{train_config.model_name}/results/saliency_{args.epoch}"
+ f"_{data_set_name}.pkl",
"wb",
) as f:
pickle.dump(map_dict, f)
# logger.info("Calculate gradients for inputs")
# # Cut off last layer of the model for saliency maps
# cutted_model = model.layers[-1].output
# # Define the last node for the discriminant output
# disc = Lambda(getDiscriminant, output_shape=discriminant_output_shape)(
# cutted_model
# )
# # Define the computation graph for the model
# model = Model(model.inputs, disc)
# # Define boolean mask to filter placeholder tracks
# boolMask = (np.sum(X_test_trk, axis=-1) != 0).astype(bool)
# # Define the number of true tracks per jet as a mask
# nTrks = np.sum(boolMask, axis=-1)
# # Get score for the dips prediction
# Db = GetScore(
# pb=pred_dips[:, b_index],
# pc=pred_dips[:, c_index],
# pu=pred_dips[:, u_index],
# )
# # Init small dict
# map_dict = {}
# # Iterate over different beff, jet flavours and passed options
# for target_beff in [60, 70, 77, 85]:
# for jet_flavour in [0, 1, 2]:
# for PassBool in [True, False]:
# # Get the Db value for a specific flavour
# Db_flavour = Db[Y_test[:, jet_flavour].astype(bool)]
# # Get the cutvalue for the specific WP
# cutvalue = np.percentile(Db_flavour, (100 - target_beff))
# # Set PassBool masking
# if PassBool is True:
# mask = Y_test[:, jet_flavour].astype(bool)
# mask = mask & (nTrks == 8)
# mask = mask & (Db > cutvalue)
# elif PassBool is False:
# mask = Y_test[:, jet_flavour].astype(bool)
# mask = mask & (nTrks == 8)
# mask = mask & (Db < cutvalue)
# # Get gradient map
# gradient_map = get_gradients(model, X_test_trk[mask], 15000)
# # Turn gradient map for plotting
# gradient_map = np.swapaxes(gradient_map, 1, 2)
# # Mean over the jets
# gradient_map = np.mean(gradient_map, axis=0)
# map_dict.update(
# {f"{target_beff}_{jet_flavour}_{PassBool}": gradient_map}
# )
# # Create results dir and .h5 ile
# os.system(f"mkdir -p {train_config.model_name}/results")
# with open(
# f"{train_config.model_name}/results/saliency_{args.epoch}"
# + f"_{data_set_name}.pkl",
# "wb",
# ) as f:
# pickle.dump(map_dict, f)
def EvaluateModelDL1(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment