Commit 6952e5ee authored by Manuel Guth's avatar Manuel Guth
Browse files

Merge branch alfroch-adding-model-save with refs/heads/master into refs/merge-requests/379/train

parents 57d2385d 56c2501a
Pipeline #3501546 passed with stages
in 10 minutes and 29 seconds
......@@ -232,6 +232,7 @@ The different options are briefly explained here:
| `exclude` | List | Necessary | List of variables that are excluded from training. Only compatible with DL1r training. To include all, just give an empty list. |
| `NN_structure` | None | Necessary | A dict where all important information for the training are defined. |
| `tagger` | String | Necessary | Name of the tagger that is used/to be trained. |
| `load_optimiser` | Bool | Optional | When loading a model (via `model_file`), you can load the optimiser state for continuing a training (`True`) or initialize a new optimiser to use the model as a start point for a fresh training (`False`). Default is `True`.
| `lr` | Float | Necessary | Learning rate which is used for training. |
| `batch_size` | Int | Necessary | Batch size which is used for training. |
| `epochs` | Int | Necessary | Number of epochs of the training. |
......
......@@ -190,6 +190,7 @@ The different options are briefly explained here:
|`tracks_name`| String| Necessary* | Name of the tracks data-set to use for training and evaluation, default is "tracks". <br />* ***This option is necessary when using tracks, but, when working with old preprpocessed files (before January 2022) this option has to be removed form the config file to ensure compatibility*** |
| `NN_structure` | None | Necessary | A dict where all important information for the training are defined. |
| `tagger` | String | Necessary | Name of the tagger that is used/to be trained. |
| `load_optimiser` | Bool | Optional | When loading a model (via `model_file`), you can load the optimiser state for continuing a training (`True`) or initialize a new optimiser to use the model as a start point for a fresh training (`False`). Default is `True`.
| `lr` | Float | Necessary | Learning rate which is used for training. |
| `batch_size` | Int | Necessary | Batch size which is used for training. |
| `epochs` | Int | Necessary | Number of epochs of the training. |
......
# Set modelname and path to Pflow preprocessing config file
model_name: dips_conditional_attention_r21
model_name: CADS_r21
preprocess_config: <path_palce_holder>/umami/examples/PFlow-Preprocessing.yaml
# Add here a pretrained model to start with.
......@@ -62,7 +62,7 @@ tracks_name: "tracks"
# Values for the neural network
NN_structure:
# Decide, which tagger is used
tagger: "dips_cond_att"
tagger: "cads"
# NN Training parameters
lr: 0.001
......
......@@ -173,7 +173,7 @@ test_train_tfrecords_dips:
- test_dips_model_tfrecords/
- coverage_files/
test_train_dips_cond_att:
test_train_cads:
<<: *test_template
stage: integration_test_tagger
needs: ["test_preprocessing_umami_count", "test_preprocessing_umami_pdf", "test_preprocessing_umami_weighting"]
......@@ -182,12 +182,12 @@ test_train_dips_cond_att:
- test_preprocessing_umami_pdf
- test_preprocessing_umami_weighting
script:
- pytest --cov=./ --cov-report= ./umami/tests/integration/test_train.py -v -s --junitxml=report.xml -k "test_train_dips_cond_att"
- cp .coverage ./coverage_files/.coverage.test_train_dips_cond_att
- pytest --cov=./ --cov-report= ./umami/tests/integration/test_train.py -v -s --junitxml=report.xml -k "test_train_cads"
- cp .coverage ./coverage_files/.coverage.test_train_cads
artifacts:
<<: *artifact_template
paths:
- test_dips_cond_att_model/
- test_cads_model/
- coverage_files/
test_train_dl1r:
......@@ -244,11 +244,11 @@ test_train_tfrecords_umami:
test_plot_input_vars:
<<: *test_template
stage: integration_test_plotting
needs: ["test_train_dips", "test_train_tfrecords_dips", "test_train_dips_cond_att", "test_train_dl1r", "test_train_umami", "test_train_tfrecords_umami"]
needs: ["test_train_dips", "test_train_tfrecords_dips", "test_train_cads", "test_train_dl1r", "test_train_umami", "test_train_tfrecords_umami"]
dependencies:
- test_train_dips
- test_train_tfrecords_dips
- test_train_dips_cond_att
- test_train_cads
- test_train_dl1r
- test_train_umami
- test_train_tfrecords_umami
......
......@@ -3,7 +3,7 @@ CI/CD:
- pipelines/
DIPS:
- umami/models/Model_Dips.py
- umami/models/Model_Dips_cond_att.py
- umami/models/Model_CADS.py
DL1:
- umami/models/Model_DL1.py
Documentation:
......
......@@ -324,7 +324,7 @@ def EvaluateModelDips(
dicts. The key will be this dataset name.
tagger : str
Name of the tagger that is to be evaluated. Can either be dips or
dips_cond_att depending which architecture is used.
CADS depending which architecture is used.
Raises
------
......@@ -379,8 +379,8 @@ def EvaluateModelDips(
model_file = utt.GetModelPath(model_name=train_config.model_name, epoch=args.epoch)
logger.info(f"Evaluating {model_file}")
# Check which test files need to be loaded depending on the DIPS version
if tagger == "dips_cond_att":
# Check which test files need to be loaded depending on the CADS version
if tagger.casefold() == "CADS".casefold():
# Load the test jets
X_test, X_test_trk, Y_test = utt.GetTestFile(
input_file=test_file,
......@@ -825,7 +825,7 @@ if __name__ == "__main__":
test_file_entry=test_file_identifier,
)
elif tagger_name in ("dips", "dips_cond_att"):
elif tagger_name.casefold() in ("dips", "cads"):
logger.info("Start evaluating DIPS with test files...")
for (
test_file_identifier,
......
"""Keras model of the DIPS tagger with conditional attention."""
"""Keras model of the CADS tagger."""
from umami.configuration import logger # isort:skip
import json
import h5py
import tensorflow as tf
from tensorflow.keras.callbacks import ModelCheckpoint
from tensorflow.keras.models import load_model
from tensorflow.keras.optimizers import Adam
import umami.tf_tools as utf
import umami.train_tools as utt
def Dips_model(train_config=None, input_shape=None):
"""Keras model definition of DIPS with conditional attention.
def Cads_model(train_config, input_shape):
"""Keras model definition of CADS.
Parameters
----------
train_config : object, optional
training config, by default None
input_shape : tuple, optional
dataset input shape, by default None
train_config : object
training config
input_shape : tuple
dataset input shape
Returns
-------
keras model
Dips with cond. attention keras model
CADS keras model
int
number of epochs
Number of epochs
"""
# Load NN Structure and training parameter from file
NN_structure = train_config.NN_structure
# Set NN options
batch_norm = NN_structure["Batch_Normalisation"]
dips = utf.Deepsets_model(
repeat_input_shape=input_shape,
num_conditions=NN_structure["N_Conditions"],
num_set_features=NN_structure["ppm_sizes"][-1],
sets_nodes=NN_structure["ppm_sizes"][:-1],
classif_nodes=NN_structure["dense_sizes"],
classif_output=len(NN_structure["class_labels"]),
pooling="attention",
attention_nodes=NN_structure["attention_sizes"],
condition_sets=NN_structure["ppm_condition"],
condition_attention=NN_structure["attention_condition"],
condition_classifier=NN_structure["dense_condition"],
shortcut_inputs=False,
sets_batch_norm=batch_norm,
classif_batch_norm=batch_norm,
activation="relu",
attention_softmax=False,
load_optimiser = (
NN_structure["load_optimiser"] if "load_optimiser" in NN_structure else True
)
# Print Dips model summary when log level lower or equal INFO level
if train_config.model_file is not None:
# Load CADS model from file
logger.info(f"Loading model from: {train_config.model_file}")
cads = load_model(
train_config.model_file,
{
"Sum": utf.Sum,
"Attention": utf.Attention,
"DeepSet": utf.DeepSet,
"AttentionPooling": utf.AttentionPooling,
"DenseNet": utf.DenseNet,
"ConditionalAttention": utf.ConditionalAttention,
"ConditionalDeepSet": utf.ConditionalDeepSet,
},
compile=load_optimiser,
)
else:
# Init a new cads/dips attention model
cads = utf.Deepsets_model(
repeat_input_shape=input_shape,
num_conditions=NN_structure["N_Conditions"],
num_set_features=NN_structure["ppm_sizes"][-1],
sets_nodes=NN_structure["ppm_sizes"][:-1],
classif_nodes=NN_structure["dense_sizes"],
classif_output=len(NN_structure["class_labels"]),
pooling="attention",
attention_nodes=NN_structure["attention_sizes"],
condition_sets=NN_structure["ppm_condition"],
condition_attention=NN_structure["attention_condition"],
condition_classifier=NN_structure["dense_condition"],
shortcut_inputs=False,
sets_batch_norm=NN_structure["Batch_Normalisation"],
classif_batch_norm=NN_structure["Batch_Normalisation"],
activation="relu",
attention_softmax=False,
)
if not load_optimiser or train_config.model_file is None:
# Set optimiser and loss
model_optimiser = Adam(learning_rate=NN_structure["lr"])
cads.compile(
loss="categorical_crossentropy",
optimizer=model_optimiser,
metrics=["accuracy"],
)
# Print CADS model summary when log level lower or equal INFO level
if logger.level <= 20:
dips.summary()
# Set optimiser and loss
model_optimizer = Adam(learning_rate=NN_structure["lr"])
dips.compile(
loss="categorical_crossentropy",
optimizer=model_optimizer,
metrics=["accuracy"],
)
return dips, NN_structure["epochs"]
cads.summary()
return cads, NN_structure["epochs"]
def DipsCondAtt(args, train_config, preprocess_config):
"""Training handling of DIPS tagger with conditonal attention.
def Cads(args, train_config, preprocess_config):
"""Training handling of CADS.
Parameters
----------
......@@ -96,8 +118,8 @@ def DipsCondAtt(args, train_config, preprocess_config):
# Print how much jets are used
logger.info(f"Number of Jets used for training: {nJets}")
# Init dips model
dips, epochs = Dips_model(train_config=train_config, input_shape=(nTrks, nFeatures))
# Init CADS model
cads, epochs = Cads_model(train_config=train_config, input_shape=(nTrks, nFeatures))
if NN_structure["use_sample_weights"]:
tensor_types = (
......@@ -129,7 +151,7 @@ def DipsCondAtt(args, train_config, preprocess_config):
# Get training set from generator
train_dataset = (
tf.data.Dataset.from_generator(
utf.dips_condition_generator(
utf.cads_generator(
train_file_path=train_config.train_file,
X_Name="X_train",
X_trk_Name=tracks_key,
......@@ -156,7 +178,7 @@ def DipsCondAtt(args, train_config, preprocess_config):
nEpochs = args.epochs
# Set ModelCheckpoint as callback
dips_mChkPt = ModelCheckpoint(
cads_mChkPt = ModelCheckpoint(
f"{train_config.model_name}/model_files" + "/model_epoch{epoch:03d}.h5",
monitor="val_loss",
verbose=True,
......@@ -225,12 +247,12 @@ def DipsCondAtt(args, train_config, preprocess_config):
)
logger.info("Start training")
history = dips.fit(
history = cads.fit(
train_dataset,
epochs=nEpochs,
# TODO: Add a representative validation dataset for training (shown in stdout)
# validation_data=validation_data,
callbacks=[dips_mChkPt, reduce_lr, my_callback],
callbacks=[cads_mChkPt, reduce_lr, my_callback],
steps_per_epoch=nJets / NN_structure["batch_size"],
use_multiprocessing=True,
workers=8,
......
......@@ -48,11 +48,14 @@ def DL1_model(train_config, input_shape, feature_connect_indices=None):
batch_norm = NN_structure["Batch_Normalisation"]
dropout = NN_structure["dropout"]
class_labels = NN_structure["class_labels"]
load_optimiser = (
NN_structure["load_optimiser"] if "load_optimiser" in NN_structure else True
)
# Load model from file if defined
if train_config.model_file is not None:
logger.info(f"Loading model from: {train_config.model_file}")
model = load_model(train_config.model_file, compile=False)
model = load_model(train_config.model_file, compile=load_optimiser)
else:
# Define input
......@@ -90,16 +93,18 @@ def DL1_model(train_config, input_shape, feature_connect_indices=None):
)(x)
model = Model(inputs=inputs, outputs=predictions)
# Compile model with given optimiser
model_optimiser = Adam(learning_rate=NN_structure["lr"])
model.compile(
loss="categorical_crossentropy",
optimizer=model_optimiser,
metrics=["accuracy"],
)
# Print DL1 model summary when log level lower or equal INFO level
if logger.level <= 20:
model.summary()
model_optimizer = Adam(learning_rate=NN_structure["lr"])
model.compile(
loss="categorical_crossentropy",
optimizer=model_optimizer,
metrics=["accuracy"],
)
return model, NN_structure["epochs"]
......
......@@ -49,11 +49,16 @@ def Dips_model(train_config=None, input_shape=None):
batch_norm = NN_structure["Batch_Normalisation"]
dropout = NN_structure["dropout"]
class_labels = NN_structure["class_labels"]
load_optimiser = (
NN_structure["load_optimiser"] if "load_optimiser" in NN_structure else True
)
if train_config.model_file is not None:
# Load DIPS model from file
logger.info(f"Loading model from: {train_config.model_file}")
dips = load_model(train_config.model_file, {"Sum": utf.Sum}, compile=False)
dips = load_model(
train_config.model_file, {"Sum": utf.Sum}, compile=load_optimiser
)
else:
logger.info("No modelfile provided! Initialize a new one!")
......@@ -108,17 +113,19 @@ def Dips_model(train_config=None, input_shape=None):
output = Dense(len(class_labels), activation="softmax", name="Jet_class")(F)
dips = Model(inputs=trk_inputs, outputs=output)
if not load_optimiser or train_config.model_file is None:
# Set optimier and loss
model_optimiser = Adam(learning_rate=NN_structure["lr"])
dips.compile(
loss="categorical_crossentropy",
optimizer=model_optimiser,
metrics=["accuracy"],
)
# Print Dips model summary when log level lower or equal INFO level
if logger.level <= 20:
dips.summary()
# Set optimier and loss
model_optimizer = Adam(learning_rate=NN_structure["lr"])
dips.compile(
loss="categorical_crossentropy",
optimizer=model_optimizer,
metrics=["accuracy"],
)
return dips, NN_structure["epochs"]
......
......@@ -57,14 +57,19 @@ def Umami_model(train_config=None, input_shape=None, njet_features=None):
batch_norm = NN_structure["Batch_Normalisation"]
dropout = NN_structure["dropout"]
class_labels = NN_structure["class_labels"]
load_optimiser = (
NN_structure["load_optimiser"] if "load_optimiser" in NN_structure else True
)
if train_config.model_file is not None:
# Load DIPS model from file
logger.info(f"Loading model from: {train_config.model_file}")
umami = load_model(train_config.model_file, {"Sum": utf.Sum}, compile=False)
umami = load_model(
train_config.model_file, {"Sum": utf.Sum}, compile=load_optimiser
)
else:
logger.info("No modelfile provided! Initialize a new one!")
logger.info("No modelfile provided! Initialise a new one!")
# Set the track input
trk_inputs = Input(shape=input_shape)
......@@ -151,19 +156,19 @@ def Umami_model(train_config=None, input_shape=None, njet_features=None):
inputs=[trk_inputs, jet_inputs], outputs=[dips_output, jet_output]
)
# Set optimier and loss
model_optimiser = Adam(learning_rate=NN_structure["lr"])
umami.compile(
loss="categorical_crossentropy",
loss_weights={"dips": NN_structure["dips_loss_weight"], "umami": 1},
optimizer=model_optimiser,
metrics=["accuracy"],
)
# Print Umami model summary when log level lower or equal INFO level
if logger.level <= 20:
umami.summary()
# Set optimier and loss
model_optimizer = Adam(learning_rate=NN_structure["lr"])
umami.compile(
loss="categorical_crossentropy",
loss_weights={"dips": NN_structure["dips_loss_weight"], "umami": 1},
optimizer=model_optimizer,
metrics=["accuracy"],
)
return umami, NN_structure["epochs"]
......
# flake8: noqa
# pylint: skip-file
from umami.models.Model_CADS import Cads
from umami.models.Model_Dips import Dips
from umami.models.Model_Dips_cond_att import DipsCondAtt
from umami.models.Model_DL1 import TrainLargeFile
from umami.models.Model_Umami import Umami
......@@ -60,7 +60,7 @@ def GetParser():
"--tagger",
type=str,
help="""Model type which is used.
You can either use 'dips', 'dips_cond_att', 'dl1' or 'umami'.""",
You can either use 'dips', 'cads', 'dl1' or 'umami'.""",
)
return parser.parse_args()
......@@ -98,7 +98,7 @@ def main(args, train_config, preprocess_config):
tagger = train_config.NN_structure["tagger"]
# Check if the tagger given is supported
if tagger in ["umami", "dl1", "dips", "dips_cond_att"]:
if tagger.casefold() in ["umami", "dl1", "dips", "cads"]:
# If dict is given, the re-calculation is skipped
if args.dict:
......@@ -143,7 +143,7 @@ def main(args, train_config, preprocess_config):
raise ValueError(
"""
You need to define a model type!\n
You can either use 'dips', 'dips_cond_att', 'dl1' or 'umami'.
You can either use 'dips', 'cads', 'dl1' or 'umami'.
"""
)
......
......@@ -23,13 +23,13 @@ test_dips:
- ci_ttbar_testing.h5
- ci_zpext_testing.h5
test_dips_cond_att:
test_cads:
data_subfolder: 'preprocessing'
model_name: 'test_dips_cond_att_model'
config: 'examples/Dips-Cond-Att-PFlow-Training-config.yaml'
model_name: 'test_cads_model'
config: 'examples/CADS-PFlow-Training-config.yaml'
preprocessing_config: 'examples/PFlow-Preprocessing.yaml'
plotting_config: 'examples/plotting_umami_config_dips.yaml'
testdir: '/tmp/umami/dips_cond_att'
testdir: '/tmp/umami/cads'
files:
- ci_ttbar_testing.h5
- ci_zpext_testing.h5
......
......@@ -21,7 +21,7 @@ def getConfiguration():
path_configuration = "umami/tests/integration/fixtures/testSetup.yaml"
with open(path_configuration, "r") as conf:
conf_setup = yaml.load(conf, Loader=yaml_loader)
for key in ["data_url", "test_dips_cond_att"]:
for key in ["data_url", "test_cads"]:
if key not in conf_setup.keys():
raise yaml.YAMLError(
f"Missing key in yaml file ({path_configuration}): {key}"
......@@ -306,11 +306,11 @@ class TestTraining(unittest.TestCase):
self.assertTrue(runTraining(config=config, tagger="DIPS"))
def test_train_dips_cond_att(self):
def test_train_cads(self):
"""Integration test of train.py for DIPS Conditional Attention script."""
config = prepareConfig(
tagger="dips_cond_att",
tagger="cads",
test_dir=self.test_dir,
preprocess_files_from="umami",
)
......
# flake8: noqa
# pylint: skip-file
from umami.tf_tools.Convert_to_Record import h5toTFRecordConverter
from umami.tf_tools.generators import (
dips_condition_generator,
cads_generator,
dips_generator,
dl1_generator,
umami_generator,
......@@ -18,6 +19,5 @@ from umami.tf_tools.layers import (
Sum,
)
from umami.tf_tools.load_tfrecord import TFRecordReader
from umami.tf_tools.Convert_to_Record import h5toTFRecordConverter
from umami.tf_tools.models import Deepsets_model
from umami.tf_tools.tools import GetLRReducer
......@@ -278,7 +278,7 @@ class umami_generator(Model_Generator):
yield {"input_1": batch_x_trk, "input_2": batch_x}, batch_y
class dips_condition_generator(Model_Generator):
class cads_generator(Model_Generator):
"""Generator class for CADS.
This class provides the a generator that loads the training dataset
......
......@@ -66,29 +66,29 @@ if __name__ == "__main__":
# Check for DIPS
# TODO: Switch to case syntax with python 3.10
if tagger_name == "dips":
if tagger_name.casefold() == "dips":
utm.Dips(
args=args,
train_config=train_config,
preprocess_config=preprocess_config,
)
elif tagger_name == "dl1":
elif tagger_name.casefold() == "dl1":
utm.TrainLargeFile(
args=args,
train_config=train_config,
preprocess_config=preprocess_config,
)
elif tagger_name == "umami":
elif tagger_name.casefold() == "umami":
utm.Umami(
args=args,
train_config=train_config,
preprocess_config=preprocess_config,
)
elif tagger_name == "dips_cond_att":
utm.DipsCondAtt(
elif tagger_name.casefold() == "cads":
utm.Cads(
args=args,
train_config=train_config,
preprocess_config=preprocess_config,
......@@ -98,6 +98,6 @@ if __name__ == "__main__":
raise ValueError(
f"""
Tagger {tagger_name} is not supported! Possible taggers are
dips, dl1, umami and dips_cond_att!
dips, dl1, umami and cads!
"""
)
......@@ -1448,7 +1448,7 @@ def calc_validation_metrics(
Raises
------
ValueError
If "tagger" is not dips, dl1, umami or dips_cond_att.
If "tagger" is not dips, dl1, umami or cads.
"""
# Get evaluation parameters and NN structure from train config
......@@ -1501,8 +1501,9 @@ def calc_validation_metrics(
# Init a results list
results = []
# TODO Change in Python 3.10
# Check tagger and load the correct val data
if tagger == "umami":
if tagger.casefold() == "umami":
data_dict = load_validation_data_umami(
train_config=train_config,
preprocess_config=preprocess_config,
......@@ -1510,7 +1511,7 @@ def calc_validation_metrics(
convert_to_tensor=False,
)
elif tagger == "dl1":