diff --git a/.gitignore b/.gitignore index 80badc4c5fa860c5e3a166dea057c3c6d5fd8d7a..70cb36b8882ac80cbfb8d09a6fea5bebc74e87cf 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ __pycache__ +.pytest_cache *.pyc *.pdf *.png diff --git a/environment.yml b/environment.yml index 1bcd59b6b6114612485c99a7fd077eea42c4e675..bb93a211fba54c20ec5015ae5be922bf51f36ef5 100644 --- a/environment.yml +++ b/environment.yml @@ -14,7 +14,9 @@ dependencies: - pip: - holidays=0.24 - prophet=1.0.1 - - pystan + - convertdate + - lunarcalendar + - pystan=2.19.1.1 - hipe4ml - hipe4ml-converter - optuna diff --git a/scripts/configs/DCM/test_config.json b/scripts/configs/DCM/test_config.json index 0f578ac10ba087d80ff6c3546ed00c645c3426ef..2a5732226b3bc0e119565ff1adf3bd848e74845c 100644 --- a/scripts/configs/DCM/test_config.json +++ b/scripts/configs/DCM/test_config.json @@ -1 +1,85 @@ -{"file_names": {"training": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree50K_DCM_trdrichrec_12agev.root", "test": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree5K_DCM_trdrichrec_12agev.root"}, "var_names": {"momentum": "Complex_p", "charge": "Complex_q", "mass2": "Complex_mass2", "pid": "Complex_pid"}, "vars_to_draw": ["Complex_mass2", "Complex_p"], "features_for_train": ["Complex_pT", "Complex_mass2", "Complex_n_hits_trd", "Complex_e_loss_0", "Complex_e_loss_1", "Complex_e_loss_2", "Complex_e_loss_3", "Complex_vtx_chi2", "Complex_chi2_ov_ndf_vtx", "Complex_chi2_ov_ndf_trd", "Complex_M"], "cuts": {"Complex_mass2": {"lower": -1.0, "upper": 2.0}, "Complex_pT": {"lower": 0.0, "upper": 2.0}, "Complex_p": {"lower": -12.0, "upper": 12.0}, "Complex_eta": {"lower": 0.0, "upper": 6.0}}, "hyper_params": {"values": {"n_estimators": 776, "max_depth": 7, "learning_rate": 0.08112714035375232, "gamma": 9, "subsample": 0.71582, "alpha": 17}, "ranges": {"n_estimators": [300, 1200], "max_depth": [2, 8], "learning_rate": [0.01, 0.1], "gamma": [0, 10], "subsample": [0.3, 0.9], "alpha": [0, 20]}}, "bins": [0.38661597183346746, 1.4825726561740278, 2.0401938866642655, 2.7015859877268475, 3.7153128963006865, 12.0]} \ No newline at end of file +{ + "file_names": { + "training": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree50K_DCM_trdrichrec_12agev.root", + "test": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree5K_DCM_trdrichrec_12agev.root" + }, + "var_names": { + "momentum": "Complex_p", + "charge": "Complex_q", + "mass2": "Complex_mass2", + "pid": "Complex_pid" + }, + "vars_to_draw": [ + "Complex_mass2", + "Complex_p" + ], + "features_for_train": [ + "Complex_pT", + "Complex_mass2", + "Complex_vtx_chi2", + "Complex_chi2_ov_ndf_vtx", + "Complex_M" + ], + "cuts": { + "Complex_mass2": { + "lower": -1.0, + "upper": 2.0 + }, + "Complex_pT": { + "lower": 0.0, + "upper": 2.0 + }, + "Complex_p": { + "lower": -12.0, + "upper": 12.0 + }, + "Complex_eta": { + "lower": 0.0, + "upper": 6.0 + } + }, + "hyper_params": { + "values": { + "n_estimators": 776, + "max_depth": 7, + "learning_rate": 0.08112714035375232, + "gamma": 9, + "subsample": 0.71582, + "alpha": 17 + }, + "ranges": { + "n_estimators": [ + 300, + 1200 + ], + "max_depth": [ + 2, + 8 + ], + "learning_rate": [ + 0.01, + 0.1 + ], + "gamma": [ + 0, + 10 + ], + "subsample": [ + 0.3, + 0.9 + ], + "alpha": [ + 0, + 20 + ] + } + }, + "bins": [ + 0.38661597183346746, + 1.4825726561740278, + 2.0401938866642655, + 2.7015859877268475, + 3.7153128963006865, + 12.0 + ] +} \ No newline at end of file diff --git a/scripts/jobs/train_no_plots_job.sh b/scripts/jobs/train_no_plots_job.sh new file mode 100755 index 0000000000000000000000000000000000000000..625c54465eeaec2bc32a76ed9218fccfea2f6685 --- /dev/null +++ b/scripts/jobs/train_no_plots_job.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +RESULT_DIR=$1 +NBINS=$2 +CONFIG=$3 + +eval "$(conda shell.bash hook)" +conda activate cbm23 + +#needed for pritning graphs as slurm doesn't find it automatically +export FONTCONFIG_FILE=$CONDA_PREFIX/etc/fonts/fonts.conf +export FONTCONFIG_PATH=$CONDA_PREFIX/etc/fonts/ + +cd $RESULT_DIR +INDEX=${SLURM_ARRAY_TASK_ID} +MLPID_SRC_DIR=/lustre/cbm/users/$USER/pid_ml/src +python $MLPID_SRC_DIR/train_model.py -c=$CONFIG --autobins=$NBINS >&training_output_${INDEX}.txt + diff --git a/scripts/onnx_launch.sh b/scripts/onnx_launch.sh new file mode 100755 index 0000000000000000000000000000000000000000..5f5be8e889bdd2dda9e68df0228bbba7d4cc8612 --- /dev/null +++ b/scripts/onnx_launch.sh @@ -0,0 +1,35 @@ + +#!/bin/bash +CONFIG=$PWD/configs/DCM/test_config.json +NBINS=5 + +timestamp=$(date +"%Y%m%d_%H%M%S") +NEW_DIR_NAME="onnx_${timestamp}" +RESULT_DIR=/lustre/cbm/users/$USER/pid/$NEW_DIR_NAME +mkdir -p $RESULT_DIR +echo "created directory $RESULT_DIR" +LOG_DIR=$RESULT_DIR/log +mkdir -p $LOG_DIR +mkdir -p $LOG_DIR/out +mkdir -p $LOG_DIR/error + +echo "logs can be found at $LOG_DIR" + +MLPID_SRC_DIR=/lustre/cbm/users/$USER/pid_ml/ + +# sbatch --job-name="autobin"\ +# --partition main\ +# --mem=32000\ +# --output=$LOG_DIR/out/%j.out.log \ +# --error=$LOG_DIR/error/%j.err.log \ +# --wait\ +# -- $PWD/jobs/autobin_job.sh $RESULT_DIR $NBINS $CONFIG + +sbatch --job-name="train-all" \ + -t 6:00:00 \ + --partition main\ + --mem=32000\ + --output=$LOG_DIR/out/%j.out.log \ + --error=$LOG_DIR/error/%j.err.log \ + --array=1-$NBINS\ + -- $PWD/jobs/train_no_plots_job.sh $RESULT_DIR $NBINS $CONFIG \ No newline at end of file diff --git a/src/onnx_config_generator.py b/src/onnx_config_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..18603856e89bcdfd46397a84b495c0399f7633c3 --- /dev/null +++ b/src/onnx_config_generator.py @@ -0,0 +1,86 @@ +import argparse +import os +import sys +from typing import List + +from util.json_tools import load_json_file, save_json_file + + +class OnnxConfigGenerator: + def __init__(self, path_to_config: str, output_dir: str, anti_particles: bool = False): + self.config: dict = load_json_file(path_to_config) + self.output_dir = output_dir + self.anti_particles = anti_particles + + os.chdir(output_dir) + self.bins = self.config["bins"] + self.model_paths_list = [] + + def find_models(self): + for i in range(0, len(self.bins) - 1): + lo, hi = self.bins[i: i+2] + model_name = self.__construct_model_name( + lo, hi, self.anti_particles) + model_path = f'{self.output_dir}/{model_name}/{model_name}.onnx' + assert (os.path.exists(model_path)) + model_dict = { + "lo": lo, + "hi": hi, + "path": model_path + } + self.model_paths_list.append(model_dict) + + @staticmethod + def __construct_model_name(lo: float, hi: float, anti_particles: bool = False) -> str: + if anti_particles: + return f"model_{lo:.1f}_{hi:.1f}_anti" + else: + return f"model_{lo:.1f}_{hi:.1f}_positive" + + def serialize(self): + output = { + "model_paths": self.model_paths_list, + "output_dir": self.output_dir, + "config_dir": self.config + } + save_json_file(output, f'{self.output_dir}/onnx_config.json') + + @staticmethod + def parse_args(args: List[str]): + parser = argparse.ArgumentParser( + prog="Onnx config generator", + description="Making config for PID ML models so that they later load correctly during inferrence", + ) + parser.add_argument( + "--config", + "-c", + required=True, + type=str, + help="Filename or path of config json file.", + ) + parser.add_argument( + "--output_dir", + "-o", + required=True, + type=str, + help="Name of folder containing all folders with output models", + ) + parser.add_argument( + "--antiparticles", + action="store_true", + help="If should train on particles instead of particles with positive charge.", + ) + return parser.parse_args(args) + + +if __name__ == "__main__": + args = OnnxConfigGenerator.parse_args(sys.argv[1:]) + + config_path = args.config + output_dir = args.output_dir + + ocg = OnnxConfigGenerator(config_path, output_dir) + print("initialized") + ocg.find_models() + print("models found") + ocg.serialize() diff --git a/src/train_model.py b/src/train_model.py index 669dfc3b9ee0fd143c2a3bbee131267c7f54b578..cdd213c727b2660c751ef86e4e41d0c1ce12f77e 100755 --- a/src/train_model.py +++ b/src/train_model.py @@ -11,6 +11,7 @@ from typing import List import json from hipe4ml.model_handler import ModelHandler +from hipe4ml_converter.h4ml_converter import H4MLConverter from sklearn.utils.class_weight import compute_sample_weight from util import json_tools, plotting_tools @@ -23,9 +24,10 @@ class TrainModel: Class for training the ml model """ - def __init__(self, model_hdl: ModelHandler, model_name: str): + def __init__(self, model_hdl: ModelHandler, model_name: str, config_file_name: str): self.model_hdl = model_hdl self.model_name = model_name + self.config_file_name = config_file_name def train_model_handler( self, train_test_data, sample_weights, model_hdl: ModelHandler = None @@ -56,7 +58,13 @@ class TrainModel: model_name = model_name or self.model_name model_hdl = model_hdl or self.model_hdl model_hdl.dump_model_handler(model_name) - print(f"\nModel saved as {model_name}") + converter = H4MLConverter(model_hdl) + features_for_train = json_tools.load_features_for_train( + self.config_file_name) + + converter.convert_model_onnx(len(features_for_train)) + converter.dump_model_onnx(f"{model_name}.onnx") + print(f"\nModel saved as {model_name} and {model_name}.onnx") def parse_args(args: List[str]) -> argparse.Namespace: @@ -223,7 +231,7 @@ if __name__ == "__main__": plotting_tools.opt_history_plot(study, save_plots) plotting_tools.opt_contour_plot(study, save_plots) # train model - train = TrainModel(model_hdl, model_name) + train = TrainModel(model_hdl, model_name, config_file_name) sample_weights = compute_sample_weight( # class_weight=None or {0: 1, 1: 3, 2: 1}, deleted for now class_weight="balanced",