diff --git a/ctlearn/conftest.py b/ctlearn/conftest.py index fa5e8733..e40c49c1 100644 --- a/ctlearn/conftest.py +++ b/ctlearn/conftest.py @@ -7,7 +7,6 @@ import numpy as np import pytest import shutil -from unittest import mock from astropy import units as u from astropy.table import Column, Table from traitlets.config.loader import Config @@ -15,9 +14,13 @@ from ctapipe.core import run_tool from ctapipe.io import write_table from ctapipe.utils import get_dataset_path -from ctlearn.tools import DLFrameWork +from ctlearn.tools.keras.train_model import TrainCTLearnKerasModel from ctlearn.utils import get_lst1_subarray_description +# TODO: ADD PyTorch here +TRAINING_TOOLS = {"Keras": TrainCTLearnKerasModel} +MODEL_FILE_FORMATS = {"Keras": "keras", "PyTorch": "pth"} + @pytest.fixture(scope="session") def gamma_simtel_path(): @@ -241,41 +244,41 @@ def ctlearn_trained_r1_mono_models(r1_gamma_file, r1_proton_file, tmp_path_facto telescope_type = "LST" # Loop over reconstruction tasks and train models for each combination ctlearn_trained_r1_mono_models = {} - with mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True): - for reco_task in ["type", "energy", "cameradirection"]: - # Output directory for trained model - output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" - - # Build command-line arguments - argv = [ - f"--signal={signal_dir}", - "--pattern-signal=*.r1.h5", - f"--output={output_dir}", - f"--reco={reco_task}", - "--TrainCTLearnModel.n_epochs=1", - "--TrainCTLearnModel.batch_size=2", - "--TrainCTLearnModel.dl1dh_reader_type=DLWaveformReader", - "--DLWaveformReader.sequence_length=5", - "--DLWaveformReader.focal_length_choice=EQUIVALENT", - ] - - # Include background only for classification task - if reco_task == "type": - argv.extend( - [ - f"--background={background_dir}", - "--pattern-background=*.r1.h5", - ] - ) - - # Run training - assert run_tool(DLFrameWork(config=config), argv=argv, cwd=tmp_path) == 0 + for reco_task in ["type", "energy", "cameradirection"]: + # Output directory for trained model + output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" + + # Build command-line arguments + argv = [ + f"--signal={signal_dir}", + "--pattern-signal=*.r1.h5", + f"--output={output_dir}", + f"--reco={reco_task}", + "--TrainCTLearnModel.n_epochs=1", + "--TrainCTLearnModel.batch_size=2", + "--TrainCTLearnModel.dl1dh_reader_type=DLWaveformReader", + "--DLWaveformReader.sequence_length=5", + "--DLWaveformReader.focal_length_choice=EQUIVALENT", + ] + + # Include background only for classification task + if reco_task == "type": + argv.extend( + [ + f"--background={background_dir}", + "--pattern-background=*.r1.h5", + "--DLWaveformReader.enforce_subarray_equality=False", + ] + ) - ctlearn_trained_r1_mono_models[f"{telescope_type}_{reco_task}"] = ( - output_dir / "ctlearn_model.keras" + # Run training tools + for framework, training_tool in TRAINING_TOOLS.items(): + assert run_tool(training_tool(config=config), argv=argv, cwd=tmp_path) == 0 + ctlearn_trained_r1_mono_models[f"{framework}_{telescope_type}_{reco_task}"] = ( + output_dir / f"ctlearn_model.{MODEL_FILE_FORMATS[framework]}" ) # Check that the trained model exists - assert ctlearn_trained_r1_mono_models[f"{telescope_type}_{reco_task}"].exists() + assert ctlearn_trained_r1_mono_models[f"{framework}_{telescope_type}_{reco_task}"].exists() return ctlearn_trained_r1_mono_models @@ -323,47 +326,45 @@ def ctlearn_trained_dl1_mono_models(dl1_gamma_file, dl1_proton_file, tmp_path_fa # Loop over telescope types and reconstruction tasks # and train models for each combination ctlearn_trained_dl1_mono_models = {} - with mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True): - for telescope_type, allowed_tels in telescope_types.items(): - for reco_task in ["type", "energy", "cameradirection"]: - # Output directory for trained model - output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" - - # Build command-line arguments - argv = [ - f"--signal={signal_dir}", - "--pattern-signal=*.dl1.h5", - f"--output={output_dir}", - f"--reco={reco_task}", - "--TrainCTLearnModel.n_epochs=1", - "--TrainCTLearnModel.batch_size=2", - "--DLImageReader.focal_length_choice=EQUIVALENT", - f"--DLImageReader.allowed_tels={allowed_tels}", - ] - - # Include background only for classification task - if reco_task == "type": - argv.extend( - [ - f"--background={background_dir}", - "--pattern-background=*.dl1.h5", - f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", - ] - ) - - # Run training - assert ( - run_tool(DLFrameWork(config=config), argv=argv, cwd=tmp_path) == 0 - ) + for telescope_type, allowed_tels in telescope_types.items(): + for reco_task in ["type", "energy", "cameradirection"]: + # Output directory for trained model + output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" + + # Build command-line arguments + argv = [ + f"--signal={signal_dir}", + "--pattern-signal=*.dl1.h5", + f"--output={output_dir}", + f"--reco={reco_task}", + "--TrainCTLearnModel.n_epochs=1", + "--TrainCTLearnModel.batch_size=2", + "--DLImageReader.focal_length_choice=EQUIVALENT", + f"--DLImageReader.allowed_tels={allowed_tels}", + ] - ctlearn_trained_dl1_mono_models[f"{telescope_type}_{reco_task}"] = ( - output_dir / "ctlearn_model.keras" + # Include background only for classification task + if reco_task == "type": + argv.extend( + [ + f"--background={background_dir}", + "--pattern-background=*.dl1.h5", + "--DLImageReader.enforce_subarray_equality=False", + f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", + ] ) - # Check that the trained model exists - assert ctlearn_trained_dl1_mono_models[ - f"{telescope_type}_{reco_task}" - ].exists() - return ctlearn_trained_dl1_mono_models + + # Run training tools + for framework, training_tool in TRAINING_TOOLS.items(): + assert run_tool(training_tool(config=config), argv=argv, cwd=tmp_path) == 0 + ctlearn_trained_dl1_mono_models[f"{framework}_{telescope_type}_{reco_task}"] = ( + output_dir / f"ctlearn_model.{MODEL_FILE_FORMATS[framework]}" + ) + # Check that the trained model exists + assert ctlearn_trained_dl1_mono_models[ + f"{framework}_{telescope_type}_{reco_task}" + ].exists() + return ctlearn_trained_dl1_mono_models @pytest.fixture(scope="session") @@ -404,42 +405,42 @@ def ctlearn_trained_dl1_stereo_models( # Loop over reconstruction tasks and train models for each combination ctlearn_trained_dl1_stereo_models = {} - with mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True): - for reco_task in ["type", "energy", "skydirection"]: - # Output directory for trained model - output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" - - # Build command-line arguments - argv = [ - f"--signal={signal_dir}", - "--pattern-signal=*.dl1.h5", - f"--output={output_dir}", - f"--reco={reco_task}", - "--TrainCTLearnModel.n_epochs=1", - "--TrainCTLearnModel.batch_size=2", - "--TrainCTLearnModel.stack_telescope_images=True", - "--DLImageReader.mode=stereo", - "--DLImageReader.focal_length_choice=EQUIVALENT", - f"--DLImageReader.allowed_tels={allowed_tels}", - ] - - # Include background only for classification task - if reco_task == "type": - argv.extend( - [ - f"--background={background_dir}", - "--pattern-background=*.dl1.h5", - ] - ) - - # Run training - assert run_tool(DLFrameWork(config=config), argv=argv, cwd=tmp_path) == 0 + for reco_task in ["type", "energy", "skydirection"]: + # Output directory for trained model + output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" + + # Build command-line arguments + argv = [ + f"--signal={signal_dir}", + "--pattern-signal=*.dl1.h5", + f"--output={output_dir}", + f"--reco={reco_task}", + "--TrainCTLearnModel.n_epochs=1", + "--TrainCTLearnModel.batch_size=2", + "--TrainCTLearnModel.stack_telescope_images=True", + "--DLImageReader.mode=stereo", + "--DLImageReader.focal_length_choice=EQUIVALENT", + f"--DLImageReader.allowed_tels={allowed_tels}", + ] + + # Include background only for classification task + if reco_task == "type": + argv.extend( + [ + f"--background={background_dir}", + "--pattern-background=*.dl1.h5", + "--DLImageReader.enforce_subarray_equality=False", + ] + ) - ctlearn_trained_dl1_stereo_models[f"{telescope_type}_{reco_task}"] = ( - output_dir / "ctlearn_model.keras" + # Run training tools + for framework, training_tool in TRAINING_TOOLS.items(): + assert run_tool(training_tool(config=config), argv=argv, cwd=tmp_path) == 0 + ctlearn_trained_dl1_stereo_models[f"{framework}_{telescope_type}_{reco_task}"] = ( + output_dir / f"ctlearn_model.{MODEL_FILE_FORMATS[framework]}" ) # Check that the trained model exists assert ctlearn_trained_dl1_stereo_models[ - f"{telescope_type}_{reco_task}" + f"{framework}_{telescope_type}_{reco_task}" ].exists() - return ctlearn_trained_dl1_stereo_models + return ctlearn_trained_dl1_stereo_models \ No newline at end of file diff --git a/ctlearn/core/__init__.py b/ctlearn/core/__init__.py index 996469a1..538e2120 100644 --- a/ctlearn/core/__init__.py +++ b/ctlearn/core/__init__.py @@ -1,2 +1,6 @@ -"""ctlearn command line tools. """ +ctlearn core functionalities +""" + +import ctlearn.core.keras.model # noqa: F401 +#import ctlearn.core.pytorch.model # noqa: F401 \ No newline at end of file diff --git a/ctlearn/core/keras/loader.py b/ctlearn/core/keras/loader.py new file mode 100644 index 00000000..6ddd0717 --- /dev/null +++ b/ctlearn/core/keras/loader.py @@ -0,0 +1,309 @@ +import numpy as np +import astropy.units as u +import keras +from keras.utils import Sequence, to_categorical + +from dl1_data_handler.reader import ProcessType + + +class KerasDataLoader(Sequence): + """ + Generates batches for Keras application. + + DLDataLoader is a data loader class that inherits from ``~keras.utils.Sequence``. + It is designed to handle and load data for deep learning models in a batch-wise manner. + + Attributes: + ----------- + data_reader : DLDataReader + An instance of DLDataReader to read the input data. + indices : list + List of indices to specify the data to be loaded. + tasks : list + List of tasks to be performed on the data to properly set up the labels. + batch_size : int + Size of the batch to load the data. + random_seed : int, optional + Whether to shuffle the data after each epoch with a provided random seed. + + Methods: + -------- + __len__(): + Returns the number of batches per epoch. + on_epoch_end(): + Updates indices after each epoch if random seed is provided. + __getitem__(index): + Generates one batch of data using _get_mono_item(index) or _get_stereo_item(index). + _get_mono_item(index): + Generates one batch of monoscopic data. + _get_stereo_item(index): + Generates one batch of stereoscopic data. + """ + + def __init__( + self, + DLDataReader, + indices, + tasks, + batch_size=64, + random_seed=None, + sort_by_intensity=False, + stack_telescope_images=False, + **kwargs, + ): + super().__init__(**kwargs) + "Initialization" + self.DLDataReader = DLDataReader + self.indices = indices + self.tasks = tasks + self.batch_size = batch_size + self.random_seed = random_seed + self.on_epoch_end() + self.stack_telescope_images = stack_telescope_images + self.sort_by_intensity = sort_by_intensity + + # Set the input shape based on the mode of the DLDataReader + if self.DLDataReader.__class__.__name__ != "DLFeatureVectorReader": + if self.DLDataReader.mode == "mono": + self.input_shape = self.DLDataReader.input_shape + elif self.DLDataReader.mode == "stereo": + self.input_shape = self.DLDataReader.input_shape[ + list(self.DLDataReader.selected_telescopes)[0] + ] + # Reshape inputs into proper dimensions + # for the stereo analysis with stacked images + if self.stack_telescope_images: + self.input_shape = ( + self.input_shape[1], + self.input_shape[2], + self.input_shape[0] * self.input_shape[3], + ) + + def __len__(self): + """ + Returns the number of batches per epoch. + + This method calculates the number of batches required to cover the entire dataset + based on the batch size. + + Returns: + -------- + int + Number of batches per epoch. + """ + return int(np.floor(len(self.indices) / self.batch_size)) + + def on_epoch_end(self): + """ + Updates indices after each epoch. If a random seed is provided, the indices are shuffled. + + This method is called at the end of each epoch to ensure that the data is shuffled + if the shuffle attribute is set to True. This helps in improving the training process + by providing the model with a different order of data in each epoch. + """ + if self.random_seed is not None: + np.random.seed(self.random_seed) + np.random.shuffle(self.indices) + + def __getitem__(self, index): + """ + Generate one batch of data and retrieve the features and labels. + + This method is called to generate one batch of monoscopic and stereoscopic data based on + the index provided. It calls either _get_mono_item(batch) or _get_stereo_item(batch) + based on the mode of the DLDataReader. + + Parameters: + ----------- + index : int + Index of the batch to generate. + + Returns: + -------- + tuple + A tuple containing the input data as features and the corresponding labels. + """ + # Generate indices of the batch + batch_indices = self.indices[ + index * self.batch_size : (index + 1) * self.batch_size + ] + features, labels = None, None + if self.DLDataReader.mode == "mono": + batch = self.DLDataReader.generate_mono_batch(batch_indices) + features, labels = self._get_mono_item(batch) + elif self.DLDataReader.mode == "stereo": + batch = self.DLDataReader.generate_stereo_batch(batch_indices) + features, labels = self._get_stereo_item(batch) + return features, labels + + def _get_mono_item(self, batch): + """ + Retrieve the features and labels for one batch of monoscopic data. + + This method is called to retrieve the features and labels for one batch of + monoscopic data. The labels are set up based on the tasks specified. + + Parameters: + ----------- + batch : astropy.table.Table + A table containing the data for the batch. + + Returns: + -------- + tuple + A tuple containing the input data as features and the corresponding labels. + """ + # Retrieve the telescope images and store in the features dictionary + labels = {} + features = batch["features"].data + if "type" in self.tasks: + labels["type"] = to_categorical( + batch["true_shower_primary_class"].data, + num_classes=2, + ) + # Temp fix till keras support class weights for multiple outputs or I wrote custom loss + # https://github.com/keras-team/keras/issues/11735 + if len(self.tasks) == 1: + labels = to_categorical( + batch["true_shower_primary_class"].data, + num_classes=2, + ) + if "energy" in self.tasks: + labels["energy"] = batch["log_true_energy"].data + if "skydirection" in self.tasks: + labels["skydirection"] = np.stack( + ( + batch["fov_lon"].data, + batch["fov_lat"].data, + ), + axis=1, + ) + if "cameradirection" in self.tasks: + labels["cameradirection"] = np.stack( + ( + batch["cam_coord_offset_x"].data, + batch["cam_coord_offset_y"].data, + ), + axis=1, + ) + return features, labels + + def _get_stereo_item(self, batch): + """ + Retrieve the features and labels for one batch of stereoscopic data. + + This method is called to retrieve the features and labels for one batch of + stereoscopic data. The original batch is grouped to retrieve the telescope + data for each event and then the telescope images or waveforms are stored + by the hillas intensity or stacked if required. Feature vectors can also + be retrieved if available for ``telescope``- and ``subarray``level. The + labels are set up based on the tasks specified. + + Parameters: + ----------- + batch : astropy.table.Table + A table containing the data for the batch. + + Returns: + -------- + tuple + A tuple containing the input data as features and the corresponding labels. + """ + labels = {} + if self.DLDataReader.process_type == ProcessType.Simulation: + batch_grouped = batch.group_by( + ["obs_id", "event_id", "tel_type_id", "true_shower_primary_class"] + ) + elif self.DLDataReader.process_type == ProcessType.Observation: + batch_grouped = batch.group_by(["obs_id", "event_id", "tel_type_id"]) + features, mono_feature_vectors, stereo_feature_vectors = [], [], [] + true_shower_primary_class = [] + log_true_energy = [] + fov_lon, fov_lat, angular_separation = [], [], [] + cam_coord_offset_x, cam_coord_offset_y, cam_coord_distance = [], [], [] + for group_element in batch_grouped.groups: + if "features" in batch.colnames: + if self.sort_by_intensity: + # Sort images by the hillas intensity in a given batch if requested + group_element.sort(["hillas_intensity"], reverse=True) + # Stack the telescope images for stereo analysis + if self.stack_telescope_images: + # Retrieve the telescope images + plain_features = group_element["features"].data + # Stack the telescope images along the last axis + stacked_features = np.concatenate( + [plain_features[i] for i in range(plain_features.shape[0])], + axis=-1, + ) + # Append the stacked images to the features list + # shape: (batch_size, image_shape, image_shape, n_channels * n_tel) + features.append(stacked_features) + else: + # Append the plain images to the features list + # shape: (batch_size, n_tel, image_shape, image_shape, n_channels) + features.append(group_element["features"].data) + # Retrieve the feature vectors + if "mono_feature_vectors" in batch.colnames: + mono_feature_vectors.append(group_element["mono_feature_vectors"].data) + if "stereo_feature_vectors" in batch.colnames: + stereo_feature_vectors.append( + group_element["stereo_feature_vectors"].data + ) + # Retrieve the labels for the tasks + # FIXME: This won't work for divergent pointing directions + if "type" in self.tasks: + true_shower_primary_class.append( + group_element["true_shower_primary_class"].data[0] + ) + if "energy" in self.tasks: + log_true_energy.append(group_element["log_true_energy"].data[0]) + if "skydirection" in self.tasks: + fov_lon.append(group_element["fov_lon"].data[0]) + fov_lat.append( + group_element["fov_lat"].data[0] + ) + if "cameradirection" in self.tasks: + cam_coord_offset_x.append(group_element["cam_coord_offset_x"].data) + cam_coord_offset_y.append( + group_element["cam_coord_offset_y"].data + ) + # Store the labels in the labels dictionary + if "type" in self.tasks: + labels["type"] = to_categorical( + np.array(true_shower_primary_class), + num_classes=2, + ) + # Temp fix till keras support class weights for multiple outputs or I wrote custom loss + # https://github.com/keras-team/keras/issues/11735 + if len(self.tasks) == 1: + labels = to_categorical( + np.array(true_shower_primary_class), + num_classes=2, + ) + if "energy" in self.tasks: + labels["energy"] = np.array(log_true_energy) + if "skydirection" in self.tasks: + labels["skydirection"] = np.stack( + ( + np.array(fov_lon), + np.array(fov_lat), + ), + axis=1, + ) + if "cameradirection" in self.tasks: + labels["cameradirection"] = np.stack( + ( + np.array(cam_coord_offset_x), + np.array(cam_coord_offset_y), + ), + axis=1, + ) + # Store the fatures in the features dictionary + if "features" in batch.colnames: + features = np.array(features) + # TDOO: Add support for both feature vectors + if "mono_feature_vectors" in batch.colnames: + features = np.array(mono_feature_vectors) + if "stereo_feature_vectors" in batch.colnames: + features = np.array(stereo_feature_vectors) + return features, labels diff --git a/ctlearn/tools/__init__.py b/ctlearn/tools/__init__.py index c6d237fc..d05c5556 100644 --- a/ctlearn/tools/__init__.py +++ b/ctlearn/tools/__init__.py @@ -1,37 +1,11 @@ """ctlearn command line tools. """ -import sys -import os -import warnings - -is_debug = '--debug' in sys.argv or any(arg.startswith('--log-level=DEBUG') for arg in sys.argv) -if not is_debug: - os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' - os.environ['NCCL_DEBUG'] = 'WARN' - warnings.filterwarnings("ignore", ".*NoneDefaultNotAllowedWarning.*") - warnings.filterwarnings("ignore", ".*MergeConflictWarning.*") - warnings.filterwarnings("ignore", ".*'ctlearn.tools.train_model' found in sys.modules.*") - -from .train_model import DLFrameWork -try: - from .predict_LST1 import LST1PredictionTool -except ImportError: - pass -try: - from .predict_model import MonoPredictCTLearnModel, StereoPredictCTLearnModel -except ImportError: - pass +from ctlearn.tools.predict_LST1 import LST1PredictionTool +from ctlearn.tools.predict_model import MonoPredictCTLearnModel, StereoPredictCTLearnModel __all__ = [ - "DLFrameWork", -] -try: - __all__.append("MonoPredictCTLearnModel") - __all__.append("StereoPredictCTLearnModel") -except NameError: - pass -try: - __all__.append("LST1PredictionTool") -except NameError: - pass \ No newline at end of file + "LST1PredictionTool", + "MonoPredictCTLearnModel", + "StereoPredictCTLearnModel", +] \ No newline at end of file diff --git a/ctlearn/tools/train/keras/__init__.py b/ctlearn/tools/keras/__init__.py similarity index 100% rename from ctlearn/tools/train/keras/__init__.py rename to ctlearn/tools/keras/__init__.py diff --git a/ctlearn/tools/train/keras/train_keras_model.py b/ctlearn/tools/keras/train_model.py similarity index 59% rename from ctlearn/tools/train/keras/train_keras_model.py rename to ctlearn/tools/keras/train_model.py index af09b2fe..57025c12 100644 --- a/ctlearn/tools/train/keras/train_keras_model.py +++ b/ctlearn/tools/keras/train_model.py @@ -1,65 +1,41 @@ +""" +Tool to train a Keras-based ``CTLearnModel``on R1/DL1a data using the ``DLDataReader`` and ``DLDataLoader``. +""" + import atexit -import pandas as pd -import numpy as np -import shutil import tensorflow as tf -from tqdm import tqdm -from time import time -from keras.callbacks import Callback +import keras from ctapipe.core.traits import ( Bool, CaselessStrEnum, - Path, - Float, - Int, - List, Dict, - classes_with_traits, - ComponentName, - Unicode, ) -from ctlearn.core.data_loader.loader import DLDataLoader -from ctlearn.tools.train.base_train_model import TrainCTLearnModel +from ctlearn.core.keras.loader import KerasDataLoader +from ctlearn.tools.train_model import TrainCTLearnModel from ctlearn.core.model import CTLearnModel from ctlearn.utils import validate_trait_dict -try: - import keras -except ImportError: - raise ImportError("keras is not installed in your environment!") - -class TqdmProgressBar(Callback): - def on_epoch_begin(self, epoch, logs=None): - self.epoch_start_time = time() - self.progress_bar = tqdm(total=self.params['steps'], desc=f'Epoc {epoch + 1}/{self.params["epochs"]}', unit='batch') - - def on_epoch_end(self, epoch, logs=None): - epoch_duration = time() - self.epoch_start_time - print(f'\nEpoch Time {epoch + 1}: {epoch_duration:.2f} seconds') - self.progress_bar.close() - - def on_batch_begin(self, batch, logs=None): - self.batch_start_time = time() - def on_batch_end(self, batch, logs=None): - # batch_duration = time() - self.batch_start_time - self.progress_bar.set_postfix(loss=logs.get('loss'), val_loss=logs.get('val_loss')) - self.progress_bar.update(1) - -class TrainKerasModel(TrainCTLearnModel): +class TrainCTLearnKerasModel(TrainCTLearnModel): """ - Tool to train a ``~ctlearn.core.model.CTLearnModel`` on R1/DL1a data using keras. + Tool to train a ``~ctlearn.core.model.CTLearnModel`` Keras-based model on R1/DL1a data. - The tool sets up the keras model using the specified optimizer and callbacks. The keras model is trained - on the input data (R1 calibrated waveforms or DL1a images) and saved in the output directory. + The tool trains a CTLearn Keras-based model on the input data (R1 calibrated waveforms or DL1a images) and + saves the trained model in the output directory. The input data is loaded from the input directories + for signal and background events using the ``~dl1_data_handler.reader.DLDataReader`` and + ``~dl1_data_handler.loader.DLDataLoader``. The tool supports the following reconstruction tasks: + - Classification of the primary particle type (gamma/proton) + - Regression of the primary particle energy + - Regression of the primary particle arrival direction based on the offsets in camera coordinates + - Regression of the primary particle arrival direction based on the offsets in sky coordinates """ name = "ctlearn-train-keras-model" description = __doc__ examples = """ - To train a CTLearn model for the classification of the primary particle type: + To train a Keras-based CTLearn model for the classification of the primary particle type: > ctlearn-train-keras-model \\ --signal /path/to/your/gammas_dl1_dir/ \\ --pattern-signal "gamma_*_run1.dl1.h5" \\ @@ -70,15 +46,16 @@ class TrainKerasModel(TrainCTLearnModel): --output /path/to/your/type/ \\ --reco type \\ - To train a CTLearn model for the regression of the primary particle energy: + To train a Keras-based CTLearn model for the regression of the primary particle energy: > ctlearn-train-keras-model \\ + --TrainCTLearnModel.model_type=KerasResNet \\ --signal /path/to/your/gammas_dl1_dir/ \\ --pattern-signal "gamma_*_run1.dl1.h5" \\ --pattern-signal "gamma_*_run10.dl1.h5" \\ --output /path/to/your/energy/ \\ --reco energy \\ - To train a CTLearn model for the regression of the primary particle + To train a Keras-based CTLearn model for the regression of the primary particle arrival direction based on the offsets in camera coordinates: > ctlearn-train-keras-model \\ --signal /path/to/your/gammas_dl1_dir/ \\ @@ -87,7 +64,7 @@ class TrainKerasModel(TrainCTLearnModel): --output /path/to/your/direction/ \\ --reco cameradirection \\ - To train a CTLearn model for the regression of the primary particle + To train a Keras-based CTLearn model for the regression of the primary particle arrival direction based on the offsets in sky coordinates: > ctlearn-train-keras-model \\ --signal /path/to/your/gammas_dl1_dir/ \\ @@ -97,9 +74,6 @@ class TrainKerasModel(TrainCTLearnModel): --reco skydirection \\ """ - model_type = ComponentName( - CTLearnModel, default_value="KerasResNet" - ).tag(config=True) save_best_validation_only = Bool( default_value=True, @@ -125,73 +99,38 @@ class TrainKerasModel(TrainCTLearnModel): ) ).tag(config=True) - aliases = { - **TrainCTLearnModel.aliases, - } - - def setup(self): - - print(tf.config.list_physical_devices('GPU')) - # Create a MirroredStrategy. - self._keras_strategy = tf.distribute.MirroredStrategy() - atexit.register(self._keras_strategy._extended._collective_ops._lock.locked) # type: ignore - self.log.info("Number of devices: %s", self._keras_strategy.num_replicas_in_sync) - # print(self.framework_type) - super().setup() - - # Set up the data loaders for training and validation - indices = list(range(self.dl1dh_reader._get_n_events())) - # Shuffle the indices before the training/validation split - np.random.seed(self.random_seed) - np.random.shuffle(indices) - n_validation_examples = int( - self.validation_split * self.dl1dh_reader._get_n_events() - ) - training_indices = indices[n_validation_examples:] - validation_indices = indices[:n_validation_examples] - # Set self.strategy.num_replicas_in_sync to 1 in case that does not exist (Pytorch) - if not hasattr(self, "_keras_strategy"): - self._keras_strategy = type("FakeStrategy", (), {"num_replicas_in_sync": 1})() - print("num_replicas_in_sync:", self._keras_strategy.num_replicas_in_sync) + def setup_framework(self): + # Create a MirroredStrategy. + self.strategy = tf.distribute.MirroredStrategy() + atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore + self.log.info("Number of devices: %s", self.strategy.num_replicas_in_sync) - print("BASE TRAIN FRAMEWORK", self.framework_type) - - self.training_loader = DLDataLoader.create( - framework=self.framework_type, - DLDataReader=self.dl1dh_reader, - indices=training_indices, + # Init the DLDataLoader for the + self.training_loader = KerasDataLoader( + self.dl1dh_reader, + self.training_indices, tasks=self.reco_tasks, - batch_size=self.batch_size * self._keras_strategy.num_replicas_in_sync, + batch_size=self.batch_size * self.strategy.num_replicas_in_sync, random_seed=self.random_seed, sort_by_intensity=self.sort_by_intensity, stack_telescope_images=self.stack_telescope_images, ) - - self.validation_loader = DLDataLoader.create( - framework=self.framework_type, - DLDataReader=self.dl1dh_reader, - indices=validation_indices, + self.validation_loader = KerasDataLoader( + self.dl1dh_reader, + self.validation_indices, tasks=self.reco_tasks, - batch_size=self.batch_size * self._keras_strategy.num_replicas_in_sync, + batch_size=self.batch_size * self.strategy.num_replicas_in_sync, random_seed=self.random_seed, sort_by_intensity=self.sort_by_intensity, stack_telescope_images=self.stack_telescope_images, ) - - def start(self): - print("Start KERAS") - print("EPOCHS:",self.n_epochs) - print("save_onnx:",self.save_onnx) - # Set up the keras callbacks + + # Set up the callbacks monitor = "val_loss" monitor_mode = "min" # Model checkpoint callback - # Temp fix for supporting keras2 & keras3 - if int(keras.__version__.split(".")[0]) >= 3: - model_path = f"{self.output_dir}/ctlearn_model.keras" - else: - model_path = f"{self.output_dir}/ctlearn_model.cpk" + model_path = f"{self.output_dir}/ctlearn_model.keras" model_checkpoint_callback = keras.callbacks.ModelCheckpoint( filepath=model_path, monitor=monitor, @@ -207,23 +146,32 @@ def start(self): csv_logger_callback = keras.callbacks.CSVLogger( filename=f"{self.output_dir}/training_log.csv", append=True ) - self.callbacks = [model_checkpoint_callback, tensorboard_callback, csv_logger_callback] - + self.callbacks = [ + model_checkpoint_callback, + tensorboard_callback, + csv_logger_callback, + ] + if self.early_stopping is not None: # EarlyStopping callback - validate_trait_dict(self.early_stopping, ["monitor", "patience", "verbose", "restore_best_weights"]) + validate_trait_dict( + self.early_stopping, + ["monitor", "patience", "verbose", "restore_best_weights"], + ) early_stopping_callback = keras.callbacks.EarlyStopping( - monitor=self.early_stopping["monitor"], - patience=self.early_stopping["patience"], - verbose=self.early_stopping["verbose"], - restore_best_weights=self.early_stopping["restore_best_weights"] + monitor=self.early_stopping["monitor"], + patience=self.early_stopping["patience"], + verbose=self.early_stopping["verbose"], + restore_best_weights=self.early_stopping["restore_best_weights"], ) self.callbacks.append(early_stopping_callback) # Learning rate reducing callback if self.lr_reducing is not None: # Validate the learning rate reducing parameters - validate_trait_dict(self.lr_reducing, ["factor", "patience", "min_delta", "min_lr"]) + validate_trait_dict( + self.lr_reducing, ["factor", "patience", "min_delta", "min_lr"] + ) lr_reducing_callback = keras.callbacks.ReduceLROnPlateau( monitor=monitor, factor=self.lr_reducing["factor"], @@ -234,12 +182,15 @@ def start(self): min_lr=self.lr_reducing["min_lr"], ) self.callbacks.append(lr_reducing_callback) + + def start(self): + # Open a strategy scope. - with self._keras_strategy.scope(): + with self.strategy.scope(): # Construct the model self.log.info("Setting up the model.") self.model = CTLearnModel.from_name( - self.model_type, + f"Keras{self.model_type}", input_shape=self.training_loader.input_shape, tasks=self.reco_tasks, parent=self, @@ -247,7 +198,7 @@ def start(self): # Validate the optimizer parameters validate_trait_dict(self.optimizer, ["name", "base_learning_rate"]) # Set the learning rate for the optimizer - learning_rate = self.optimizer["base_learning_rate"] + learning_rate = self.optimizer["base_learning_rate"] # Set the epsilon for the Adam optimizer adam_epsilon = None if self.optimizer["name"] == "Adam": @@ -266,7 +217,10 @@ def start(self): keras.optimizers.Adam, dict(learning_rate=learning_rate, epsilon=adam_epsilon), ), - "RMSProp": (keras.optimizers.RMSprop, dict(learning_rate=learning_rate)), + "RMSProp": ( + keras.optimizers.RMSprop, + dict(learning_rate=learning_rate), + ), "SGD": (keras.optimizers.SGD, dict(learning_rate=learning_rate)), } # Get the optimizer function and arguments @@ -275,10 +229,9 @@ def start(self): losses, metrics = self._get_losses_and_mertics(self.reco_tasks) # Compile the model self.log.info("Compiling CTLearn model.") - self.model.compile(optimizer=optimizer_fn(**optimizer_args), loss=losses, metrics=metrics) - - tqdm_callback = TqdmProgressBar() - self.callbacks.append(tqdm_callback) + self.model.compile( + optimizer=optimizer_fn(**optimizer_args), loss=losses, metrics=metrics + ) # Train and evaluate the model self.log.info("Training and evaluating...") @@ -292,24 +245,6 @@ def start(self): ) self.log.info("Training and evaluating finished succesfully!") - def finish(self): - # Saving model weights in onnx format - if self.save_onnx: - self.log.info("Converting Keras model into ONNX format...") - self.log.info("Make sure tf2onnx is installed in your enviroment!") - try: - import tf2onnx - except ImportError: - raise ImportError("tf2onnx is not installed in your environment!") - - output_path = f"{self.output_dir}/ctlearn_model.onnx" - tf2onnx.convert.from_keras( - self.model, input_signature=self.model.input_layer.input._type_spec, output_path=output_path - ) - self.log.info("ONNX model saved in %s", self.output_dir) - - self.log.info("Tool is shutting down") - def _get_losses_and_mertics(self, tasks): """ Build the fully connected head for the CTLearn model. @@ -353,10 +288,24 @@ def _get_losses_and_mertics(self, tasks): losses["cameradirection"] = keras.losses.MeanAbsoluteError( reduction="sum_over_batch_size" ) - metrics["cameradirection"] = keras.metrics.MeanAbsoluteError(name="mae_cameradirection") + metrics["cameradirection"] = keras.metrics.MeanAbsoluteError( + name="mae_cameradirection" + ) if "skydirection" in self.reco_tasks: losses["skydirection"] = keras.losses.MeanAbsoluteError( reduction="sum_over_batch_size" ) - metrics["skydirection"] = keras.metrics.MeanAbsoluteError(name="mae_skydirection") - return losses, metrics \ No newline at end of file + metrics["skydirection"] = keras.metrics.MeanAbsoluteError( + name="mae_skydirection" + ) + return losses, metrics + + +def main(): + # Run the tool + tool = TrainCTLearnKerasModel() + tool.run() + + +if __name__ == "main": + main() \ No newline at end of file diff --git a/ctlearn/tools/pytorch/CTLearnPL.py b/ctlearn/tools/pytorch/CTLearnPL.py new file mode 100644 index 00000000..62fea3d1 --- /dev/null +++ b/ctlearn/tools/pytorch/CTLearnPL.py @@ -0,0 +1,1783 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +import numpy as np +from ctlearn.core.pytorch.nets.loss_functions.loss_functions import ( + FocalLoss, + VectorLoss, + AngularDistance, + cosine_direction_loss, +) +from ctlearn.core.pytorch.nets.optimizer.optimizer import one_cycle +from tqdm import tqdm +import os +import math +import torch.optim.lr_scheduler as lr_scheduler +from sklearn.metrics import confusion_matrix +from torchmetrics.classification import ConfusionMatrix, MulticlassPrecision, MulticlassF1Score + +from io import BytesIO +import torch.nn.functional as F +from ctlearn.core.pytorch.utils import utils +from ctlearn.core.pytorch.visualization.vis_utils import ( + plot_confusion_matrix, + plot_energy_resolution_error, + plot_direction_resolution_error, +) + +# import utils_torch +import pytorch_lightning as pl +from torchmetrics import Accuracy +import matplotlib.pyplot as plt +from astropy.coordinates import SkyCoord +from astropy.time import Time +import astropy.units as u +from ctlearn.core.ctlearn_enum import Task, Mode +import gc +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from ctlearn.core.pytorch.nets.loss_functions.loss_functions import evidential_regression_loss +import inspect + +import matplotlib +matplotlib.use('Agg') +torch.cuda.empty_cache() +os.environ["CUDA_LAUNCH_BLOCKING"] = "1" + +class CTLearnTrainer(pl.Trainer): + def __init__(self,**kwargs): + super(CTLearnTrainer, self).__init__(**kwargs) + self.multi_gpu=False + if self.world_size > 1: + self.multi_gpu=True + + def predict( + self, + model, + dataloaders=None, + return_predictions=False, + h5_file_name="./results.r1.dl2.h5", + task: Task = None, + mode: Mode = None, + **kwargs, + ): + + # ---------------------------------------------------------------- + # Prediction + # ---------------------------------------------------------------- + self.class_predictions = [] + self.energy_predictions = [] + self.direction_predictions = [] + self.event_id_list = [] + self.obs_id_list = [] + self.labels_energy_list = [] + self.labels_direction_list = [] + self.labels_true_alt_az_list = [] + self.hillas_list = [] + self.pointing_dir = {} + self.h5_file_name = h5_file_name + self.task = task + self.mode = mode + # super(CTLearnTrainer, self).__init__(kwargs) + print("Using CustomTrainer's predict method.") + + if dataloaders is None: + raise ValueError("A dataloader must be provided for prediction.") + + model.eval() + with torch.no_grad(): + # Call your custom logic here + results = model.generate_results( + input_data_loader=dataloaders, + h5_file_name=None, + task=task, + mode=mode, + ) + + if return_predictions: + return results + + print("Prediction complete.") + + def get_log_dir(self) -> str: + return self.logger.log_dir + +class CTLearnPL(pl.LightningModule): + + # def setup(self,stage): + + # # self.dummy = None + + # self.dummy = self.occupy_free_gpu_memory(self.device) + + def occupy_free_gpu_memory(self,device): + # Get free and total GPU memory using CUDA APIs + free_mem, total_mem = torch.cuda.mem_get_info(device) + + # Leave a margin (e.g., 100 MB) + margin = 20000 * 1024 * 1024 # 1000 MB + mem_to_allocate = max(0, free_mem - margin) + + if mem_to_allocate > 0: + try: + # Each float32 element takes 4 bytes + num_elements = mem_to_allocate // 4 + dummy_tensor = torch.zeros(num_elements, dtype=torch.float32, device=device) + print(f"Occupied {mem_to_allocate/1024**2:.2f} MB of GPU memory on {device}") + return dummy_tensor + except RuntimeError as e: + print(f"Could not allocate memory: {e}") + return None + else: + print("Not enough free GPU memory to occupy.") + return None + + def __init__( + self, + model, + save_folder, + task: Task, + mode: Mode, + parameters, + train_loader=None, + val_loader=None, + test_val_loader=None, + train_dataset=None, + val_dataset=None, + num_inputs=1, + k=3, + + ): + super(CTLearnPL, self).__init__() + + self.task = task + self.mode = mode + + self.save_folder = save_folder + self.model = model + + # torch.autograd.set_detect_anomaly(True) + + self._device = torch.device(parameters["arch"]["device"]) + self.device_type = parameters["arch"]["device"] + + self.model.to(self.device) + + + + # # Ejecutar la función + # self.dummy = None + + # self.occupy_free_gpu_memory(self.device) + self.k = k # Number of top results to save + # Get the number of inputs of the net + sig = inspect.signature(model.forward) + num_inputs = len(sig.parameters) + # Detect if model is diffusive + if hasattr(self.model,"T"): + self.is_difussion=True + else: + self.is_difussion=False + print("Diffusion: ",self.is_difussion) + + self.num_inputs = num_inputs + self.train_loader = train_loader + self.val_loader = val_loader + self.train_dataset = train_dataset + self.val_dataset = val_dataset + self.test_val_loader = test_val_loader + self.class_names = ["gamma", "proton"] + self.val_proton_gammanes = [] + self.val_gamma_gammanes = [] + # Hyperparameters + self.set_hyperparameters(parameters) + + self.optimizer = None + self.scheduler = None + # self.class_weights = torch.tensor([1.0, 1.3], dtype=torch.float32).to(self.device).contiguous() + + # Loss Function + if self.task == Task.type: + weights = ( + torch.tensor([parameters['class_weight'][0],parameters['class_weight'][1]], dtype=torch.float32).to(self.device).contiguous() + ) # [1.0, 1.3] + self.register_buffer('class_weights', weights) + self.criterion_class = nn.CrossEntropyLoss( + weight=self.class_weights, reduction="mean" + ) + + + self.criterion_energy_value = torch.nn.L1Loss(reduction="mean") + self.criterion_direction = torch.nn.SmoothL1Loss() # nn.MSELoss() + self.criterion_magnitud = torch.nn.L1Loss(reduction="mean") + + self.criterion_direction_none = torch.nn.SmoothL1Loss(reduction="none") # nn.MSELoss() + self.criterion_magnitud_none = torch.nn.L1Loss(reduction="none") + self.criterion_alt_az_l1_none = torch.nn.L1Loss(reduction="none") + self.criterion_energy_value_none = torch.nn.L1Loss(reduction="none") + + self.criterion_direction = torch.nn.L1Loss(reduction="mean") # nn.MSELoss() + self.criterion_vector = VectorLoss(alpha=0.1, reduction="mean") + self.criterion_alt_az_l1 = torch.nn.L1Loss(reduction="mean") + + self.criterion_alt_az = evidential_regression_loss(lamb=0.01, reduction="mean") + + + self.best_loss = float("inf") + self.best_accuracy = 0 + # Best Metrics Tracking + self.best_losses = [(float("inf"), None)] * k + self.best_accuracies = [(0, None)] * k + self.best_validation_accuracy = 0 + + self.correct_classification = 0 + + self.f1_score_val = MulticlassF1Score( + num_classes=2, + dist_sync_on_step=True, + ) + + self.f1_score_train = MulticlassF1Score( + num_classes=2, + dist_sync_on_step=True, + ) + + self.precision_val = MulticlassPrecision(num_classes=2,dist_sync_on_step=True,) + self.precision_train = MulticlassPrecision(num_classes=2,dist_sync_on_step=True,) + + self.class_train_accuracy = Accuracy( + task="multiclass", + num_classes=parameters["model"]["model_type"]["parameters"]["num_outputs"], + dist_sync_on_step=True # GPUs Sync + ) + + self.class_val_accuracy = Accuracy( + task="multiclass", + num_classes=parameters["model"]["model_type"]["parameters"]["num_outputs"], + dist_sync_on_step=True # GPUs Sync + ) + + self.confusion_matrix = ConfusionMatrix(num_classes=2, task="multiclass",dist_sync_on_step=True) + + self.loss_train_sum = 0.0 + self.num_train_batches = 0 + + # ---------------------------------------------------------------- + # List used to plot on tensorboard + # ---------------------------------------------------------------- + + # ---------------------------------------------------------------- + # Validation + # ---------------------------------------------------------------- + self.loss_val_sum = 0.0 + self.num_val_batches = 0 + + + self.val_angular_diff_list = [] + self.val_energy_diff_list = [] + + self.val_energy_pred_list = [] + + self.val_alt_pred_list = [] + self.val_az_pred_list = [] + self.val_alt_label_list = [] + self.val_az_label_list = [] + + self.loss_val_distance = 0.0 + self.loss_val_dx_dy = 0.0 + self.loss_val_angular_error = 0.0 + + self.val_energy_label_list = [] + self.val_hillas_intensity_list = [] + + # ---------------------------------------------------------------- + # Training + # ---------------------------------------------------------------- + self.loss_train_distance = 0.0 + self.loss_train_dx_dy = 0.0 + self.loss_train_angular_error = 0.0 + + self.alt_off_list = [] + self.az_off_list = [] + + # ---------------------------------------------------------------- + # Predictions + # ---------------------------------------------------------------- + + self.predictions = [] + + self.class_predictions = [] + self.energy_predictions = [] + self.direction_predictions = [] + self.event_id_list = [] + self.obs_id_list = [] + self.labels_energy_list = [] + self.labels_direction_list = [] + self.labels_true_alt_az_list = [] + self.hillas_list = [] + self.pointing = [] + # ---------------------------------------------------------------------------------------------------------- + def train_dataloader(self): + return self.train_loader + # ---------------------------------------------------------------------------------------------------------- + def set_hyperparameters(self, parameters): + # Hyperparameters + self.learning_rate = float(parameters["hyp"]["learning_rate"]) + self.adam_epsilon = float(parameters["hyp"]["adam_epsilon"]) + self.momentum = float(parameters["hyp"]["momentum"]) + self.weight_decay = float(parameters["hyp"]["weight_decay"]) + self.num_epochs = int(parameters["hyp"]["epochs"]) + + self.start_epoch = parameters["hyp"]["start_epoch"] + self.steps_epoch = parameters["hyp"]["steps_epoch"] + self.lrf = float(parameters["hyp"]["lrf"]) + self.optimizer_type = str(parameters["hyp"]["optimizer"]).lower() + self.l2_lambda = float(parameters["hyp"]["l2_lambda"]) + # ---------------------------------------------------------------------------------------------------------- + def forward(self, x, y): + return self.model(x, y) + # ---------------------------------------------------------------------------------------------------------- + def save_checkpoint(self, save_folder, metric_value, filename_prefix, is_loss=True): + """Save checkpoint and manage top k checkpoints for loss or accuracy.""" + filename = os.path.join( + save_folder, f"{filename_prefix}_{metric_value:.16f}.pth" + ) + torch.save( + { + "model_state_dict": self.model.state_dict(), + "optimizer_state_dict": self.optimizer.state_dict(), + "metric_value": metric_value, + }, + filename, + ) + + # Determine the list to update + current_list = self.best_losses if is_loss else self.best_accuracies + current_list.append((metric_value, filename)) + current_list.sort(reverse=not is_loss) + if len(current_list) > self.k: + removed_metric, removed_file = current_list.pop() + if removed_file and os.path.exists(removed_file): + # TODO: Add a Lock to avoid problems when using multiple GPUS + try: + os.remove(removed_file) # Remove the worst performing file + except FileNotFoundError: + pass # File doesn't exist, continue execution + + + # ---------------------------------------------------------------------------------------------------------- + def compute_type_loss( + self, classification_pred, labels_class, test_val=False, training=False + ): + target = labels_class.to(torch.int64) + + # Cálculo de la loss con F.cross_entropy + loss_class = F.cross_entropy(classification_pred, target, weight=self.class_weights.to(classification_pred.device), reduction='mean') + + # Calculate accuracy + predicted = torch.softmax(classification_pred, dim=1) + predicted = predicted.argmax(dim=1) + + + accuracy = 0 + precision = 0 + loss = loss_class + + if training: + + self.class_train_accuracy.update(predicted, labels_class) + accuracy = self.class_train_accuracy.compute().item() + self.f1_score_train.update(predicted, labels_class) + self.precision_train.update(predicted, labels_class) + precision = self.precision_train.compute().item() + else: + + self.class_val_accuracy.update(predicted, labels_class) + accuracy = self.class_val_accuracy.compute().item() + self.confusion_matrix.update(predicted, labels_class) + self.f1_score_val.update(predicted, labels_class) + self.precision_val.update(predicted, labels_class) + precision = self.precision_val.compute().item() + + return loss, accuracy, predicted, precision + # ---------------------------------------------------------------------------------------------------------- + def compute_camera_direction_loss(self, direction_pred, labels_direction, training=False): + + labels_dx_dy = labels_direction[:, 0:2] + label_distance = labels_direction[:, 2] + + + if isinstance(direction_pred, tuple): + direction_pred = list(direction_pred) + + pred_dx_dy = direction_pred[0][:,0:2].unsqueeze(-1) + pred_distance = direction_pred[0][:,2].unsqueeze(-1) + else: + pred_dx_dy = direction_pred[:,0:2].unsqueeze(-1) + pred_distance = direction_pred[:,2].unsqueeze(-1) + + + loss_angular_diff = cosine_direction_loss(pred_dx_dy[:,0],pred_dx_dy[:,1], labels_dx_dy[:, 0],labels_dx_dy[:, 1]) + + + _, angular_diff = AngularDistance( + pred_dx_dy[:,0], + labels_dx_dy[:, 0], + pred_dx_dy[:,1], + labels_dx_dy[:, 1], + reduction="None", + ) + + vector_cam_distance = torch.sqrt(pred_dx_dy[:,0]**2 + pred_dx_dy[:,1]**2) + loss_dx_dy = self.criterion_alt_az_l1(pred_dx_dy, labels_dx_dy) + loss_distance = self.criterion_magnitud(label_distance, pred_distance) + loss_distance_dx_dy = self.criterion_magnitud(label_distance, vector_cam_distance) + + loss = loss_dx_dy + loss_distance + loss_distance_dx_dy + loss_angular_diff + return loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_energy_loss( + self, energy_pred, labels_energy, test_val=False, training=False + ): + + loss_energy = self.criterion_energy_value(energy_pred, labels_energy) + loss = loss_energy + #----------------------------------------- + # k=0.6 + # e_thrs= -0.3 + # loss_energy = self.criterion_energy_value_none(energy_pred, labels_energy) + # energy_tev = torch.pow(10,labels_energy) + # energy_weight=k*torch.exp(1+(np.log10(1/k)*(energy_tev-e_thrs))) + # weight_loss = energy_weight * (loss_energy) + + # loss = weight_loss.sum() + #----------------------------------------- + + if training == False: + energy_pred = pow(10, energy_pred) + labels_energy = pow(10, labels_energy) + energy_diff = torch.abs(energy_pred - labels_energy) + energy_diff = energy_diff.float().cpu().detach().numpy() + else: + energy_diff = None + + return loss, energy_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_energy_loss_diffusion(self, x_1, y,training=False, x_2 = None + # self, energy_pred, labels_energy, test_val=False, training=False + ): + + loss = 0 + # y = y.squeeze(-1) + y_embed = self.model.target_embedder(y) + k=0.6 + e_thrs= -0.3 + for t in range(self.model.T): + # Add noise to target (landmarks) + alpha_bar_t = self.model.alpha_bar[t] + noise = torch.randn_like(y_embed) + z_t = torch.sqrt(alpha_bar_t) * y_embed + torch.sqrt(1 - alpha_bar_t) * noise + + if x_2 is None: + # Denoise step + z, _ = self.model.blocks[t](x_1, z_t, None) # W_embed not needed + else: + z, _ = self.model.blocks[t](x_1 ,x_2, z_t, None) # W_embed not needed + + + # Loss to clean target + loss_l2 = F.mse_loss(z, y_embed) + + labels_energy = y + + # Weighted by SNR difference + step_loss = 2.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + # Final step: use regression head + if t == self.model.T - 1: + + energy_pred = self.model.regress(z) + + # labels_energy_tev = torch.pow(10,labels_energy) + # energy_pred_tev = torch.pow(10,energy_pred) + + # loss_energy = self.criterion_energy_value(energy_pred_tev, labels_energy_tev) + + loss_energy = self.criterion_energy_value(energy_pred, labels_energy) + # loss_energy = F.mse_loss(energy_pred, labels_energy,reduction="sum") + # loss_energy = self.criterion_energy_value_none(energy_pred, labels_energy) + + # energy_weight=k*torch.exp(1+(np.log10(1/k)*(energy_tev-e_thrs))) + # weight_loss = energy_weight * (loss_energy) + + # step_loss = step_loss + weight_loss.mean() + + if training == False: + energy_pred = pow(10, energy_pred) + labels_energy = pow(10, labels_energy) + energy_diff = torch.abs(energy_pred - labels_energy) + energy_diff = energy_diff.float().cpu().detach().numpy() + else: + energy_diff = None + + step_loss = step_loss + loss_energy + + loss = loss + step_loss + + return loss, energy_diff, energy_pred + # ---------------------------------------------------------------------------------------------------------- + def compute_camera_direction_loss_diffusion(self, x_1, y, labels_energy_value, x_2 = None): + + loss = 0 + y = y.squeeze(-1) + y_embed = self.model.target_embedder(y) + k=0.6 + e_thrs= -0.3 + + for t in range(self.model.T): + # Add noise to target (landmarks) + alpha_bar_t = self.model.alpha_bar[t] + noise = torch.randn_like(y_embed) + z_t = torch.sqrt(alpha_bar_t) * y_embed + torch.sqrt(1 - alpha_bar_t) * noise + + if x_2 is None: + # Denoise step + z, _ = self.model.blocks[t](x_1, z_t, None) # W_embed not needed + else: + z, _ = self.model.blocks[t](x_1,x_2, z_t, None) # W_embed not needed + + preds = self.model.regress(z) + # Loss to clean target + loss_l2 = F.mse_loss(preds, y) + + # Weighted by SNR difference + step_loss = 2.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + # Final step: use regression head + if t == self.model.T - 1: + labels_dx_dy = y[:, 0:2] + label_distance = y[:, 2] + direction_pred = self.model.regress(z) + if isinstance(direction_pred, tuple): + # Not Tested + direction_pred = list(direction_pred) + + pred_dx_dy = direction_pred[0][:,0:2].unsqueeze(-1) + pred_distance = direction_pred[0][:,2].unsqueeze(-1) + else: + pred_dx_dy = direction_pred[:,0:2] + pred_distance = direction_pred[:,2] + + # loss_angular_diff = cosine_direction_loss(pred_dx_dy[:,0],pred_dx_dy[:,1], labels_dx_dy[:, 0],labels_dx_dy[:, 1]) + loss_angular_diff = cosine_direction_loss(pred_dx_dy[:,0],pred_dx_dy[:,1], labels_dx_dy[:, 0],labels_dx_dy[:, 1],reduction="none") + _, angular_diff = AngularDistance( + pred_dx_dy[:,0], + labels_dx_dy[:, 0], + pred_dx_dy[:,1], + labels_dx_dy[:, 1], + reduction="None", + + ) + + vector_cam_distance = torch.sqrt(pred_dx_dy[:,0]**2 + pred_dx_dy[:,1]**2) + # loss_dx_dy = self.criterion_alt_az_l1(pred_dx_dy, labels_dx_dy) + # loss_distance = self.criterion_magnitud(label_distance, pred_distance) + # loss_distance_dx_dy = self.criterion_magnitud(label_distance, vector_cam_distance) + loss_dx_dy = self.criterion_alt_az_l1_none(pred_dx_dy, labels_dx_dy).mean(dim=1) + loss_distance = self.criterion_magnitud_none(label_distance, pred_distance) + loss_distance_dx_dy = self.criterion_magnitud_none(label_distance, vector_cam_distance) + + energy = torch.pow(10,labels_energy_value) + # k=4.3 + # e_thrs= 4 + # energy_weight = k*(1/(1+torch.exp(-(1/k)*(energy-e_thrs)))) + # weight_loss = energy_weight * (loss_dx_dy + loss_distance + loss_distance_dx_dy + loss_angular_diff) + + + + energy_weight=k*torch.exp(1+(np.log10(1/k)*(energy-e_thrs))) + # weight_loss = energy_weight * (loss_dx_dy + loss_distance + loss_distance_dx_dy + loss_angular_diff) + weight_loss = energy_weight * (loss_dx_dy + loss_distance + loss_distance_dx_dy) + + loss_distance=loss_distance.mean() + loss_dx_dy = loss_dx_dy.mean() + loss_distance_dx_dy = loss_distance_dx_dy.mean() + + loss_angular_diff = loss_angular_diff.mean() + + step_loss = step_loss + weight_loss.mean() + # step_loss = step_loss + loss_dx_dy + loss_distance + loss_distance_dx_dy + loss = loss + step_loss + + return loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_sky_direction_loss_diffusion(self, x, y, labels_energy_value): + + loss = 0 + y = y.squeeze(-1) + y_embed = self.model.target_embedder(y) + + for t in range(self.model.T): + # Add noise to target (landmarks) + alpha_bar_t = self.model.alpha_bar[t] + noise = torch.randn_like(y_embed) + z_t = torch.sqrt(alpha_bar_t) * y_embed + torch.sqrt(1 - alpha_bar_t) * noise + + # Denoise step + z, _ = self.model.blocks[t](x, z_t, None) # W_embed not needed + + preds = self.model.regress(z) + # Loss to clean target + loss_l2 = F.mse_loss(preds, y) + + # Weighted by SNR difference + step_loss = 2.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + # Final step: use regression head + if t == self.model.T - 1: + + + labels_dx_dy = y[:, 0:2] + label_distance = y[:, 2] + direction_pred = self.model.regress(z) + + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + # direction_pred = list(direction_pred) + # pred_az_atl = direction_pred[0][:,0:2] + # pred_separation = direction_pred[0][:,2] + # # direction_pred[0]= direction_pred[0][:,0:2] + # else: + + # pred_az_atl = direction_pred[:, 0:2] + # pred_separation = direction_pred[:, 2] + + # labels_az_alt = labels_direction[:, 0:2] + # label_separation = labels_direction[:, 2] + # loss_separation = self.criterion_direction(pred_separation, label_separation) + + # # loss_vector = self.criterion_vector(pred_dir_cartesian, labels_direction_cartesian) + # # vect_magnitud = torch.sqrt(torch.sum(pred_dir_cartesian**2, dim=1)) + # # loss_magnitud = torch.abs(1.0-vect_magnitud).sum() + + # # alt_az = utils_torch.cartesian_to_alt_az(direction[:,0:3]) + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + # loss_alt_az = self.criterion_alt_az(direction_pred, labels_direction) + # else: + # loss_alt_az = self.criterion_alt_az_l1(pred_az_atl, labels_az_alt) + + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + + # loss_angular_error, _ = AngularDistance( + # (direction_pred[0][:, 1]), + # labels_az_alt[:, 1], + # (direction_pred[0][:, 0]), + # labels_az_alt[:, 0], + # reduction="sum", + # ) + + # else: + + # loss_angular_error, _ = AngularDistance( + # (direction_pred[:, 1]), + # labels_az_alt[:, 1], + # (direction_pred[:, 0]), + # labels_az_alt[:, 0], + # reduction="sum", + # ) + + # if training == False: + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + # _, angular_diff = AngularDistance( + # (direction_pred[0][:, 1]), + # labels_az_alt[:, 0], + # (direction_pred[0][:, 0]), + # labels_az_alt[:, 1], + # reduction=None, + # ) + # else: + # _, angular_diff = AngularDistance( + # (direction_pred[:, 1]), + # labels_az_alt[:, 0], + # (direction_pred[:, 0]), + # labels_az_alt[:, 1], + # reduction=None, + # ) + # else: + # angular_diff = None + + # loss = loss_alt_az + 0.001*(loss_separation + loss_angular_error) + + # loss = loss + step_loss + + # return loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_class_weights(self, y, n_classes): + # y: tensor de etiquetas, por ejemplo torch.tensor([0,0,1,1,1,2]) + counts = torch.bincount(y, minlength=n_classes) + total = counts.sum() + # Evita división por cero: + counts = torch.clamp(counts, min=1) + weights = total / (counts * n_classes) + return weights + # ---------------------------------------------------------------------------------------------------------- + def compute_type_loss_diffusion(self, x,y, training=False): + + loss = 0 + # Get the embedding of the true label (e.g. "this is a 3") + uy = self.model.W_embed[y] + for t in range(self.model.T): + + # Get the current noise level from the schedule + alpha_bar_t = self.model.alpha_bar[t] + + # Generate random noise with the same shape as uy + noise = torch.randn_like(uy) + + # Add noise to the label embedding → this is our "noisy target" + z_t = torch.sqrt(alpha_bar_t) * uy + torch.sqrt(1 - alpha_bar_t) * noise + + # Pass the image and the noisy label through one block + z_pred, _ = self.model.blocks[t](x, z_t, self.model.W_embed) + + # Compute how far the output is from the clean label embedding + loss_l2 = F.mse_loss(z_pred, uy) + + # Weight the loss using the signal-to-noise ratio + step_loss = 0.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + + accuracy = 0 + precision = 0 + predicted = None + # If it's the final layer, add classification and KL losses + if t == self.model.T - 1: + # Get predictions from classifier + logits = self.model.classifier(z_pred) + + # class_weights = self.compute_class_weights(y, n_classes=2).to(y.device) + # loss_ce = F.cross_entropy(logits, y, class_weights) + + # Cross-entropy loss: how wrong the predicted class is + loss_ce = F.cross_entropy(logits, y) + + # KL-like loss: penalize if embedding is too far from origin + loss_kl = 0.5 * uy.pow(2).sum(dim=1).mean() + + # Add all parts together + # loss = loss + loss_ce + loss_kl + step_loss = step_loss + loss_ce + loss_kl + classification_pred,*_ = self.model(x) + + # Calculate accuracy + predicted = torch.softmax(classification_pred, dim=1) + predicted = predicted.argmax(dim=1) + + if training: + self.class_train_accuracy.update(predicted, y) + accuracy = self.class_train_accuracy.compute().item() + self.f1_score_train.update(predicted, y) + self.precision_train.update(predicted, y) + precision = self.precision_train.compute().item() + else: + self.class_val_accuracy.update(predicted, y) + accuracy = self.class_val_accuracy.compute().item() + self.confusion_matrix.update(predicted, y) + self.f1_score_val.update(predicted, y) + self.precision_val.update(predicted, y) + precision = self.precision_val.compute().item() + loss = loss + step_loss + + return loss, accuracy, predicted, precision + # ---------------------------------------------------------------------------------------------------------- + # def on_train_start(self): + # self.dummy_tensor = self.occupy_free_gpu_memory(self.device) + # ---------------------------------------------------------------------------------------------------------- + def on_train_batch_start(self, batch, batch_idx): + if self.device != torch.device("cpu"): + if batch_idx == 5 and not hasattr(self, "dummy_tensor"): + self.dummy_tensor = self.occupy_free_gpu_memory(self.device) + + # ---------------------------------------------------------------------------------------------------------- + def training_step(self, batch, batch_idx): + # ------------------------------------------------------------------ + # Read inputs (features) and labels + # ------------------------------------------------------------------ + features, labels, t = batch + loss = 0 + if len(features) > 0: + imgs = features["image"] + + if self.task == Task.type: + labels_class = labels["type"] + + labels_energy_value = labels["energy"] + if self.task == Task.energy: + labels_energy_value = labels_energy_value.to(self.device) + + + if self.task == Task.cameradirection: + labels_direction = labels["direction"] + + imgs = imgs.to(self.device) + + # Log sample images to TensorBoard every 5 epochs + if batch_idx == 0 and self.current_epoch % 5 == 0 and self.trainer.is_global_zero: + import torchvision + try: + num_samples = min(4, imgs.shape[0]) + sample_imgs = imgs[:num_samples] + + label_str = "" + if self.task == Task.type: + label_str = str(labels_class[:num_samples].tolist()) + elif self.task == Task.energy: + label_str = str(labels_energy_value[:num_samples].tolist()) + elif self.task == Task.cameradirection: + label_str = str(labels_direction[:num_samples].tolist()) + + grid = torchvision.utils.make_grid(sample_imgs, normalize=True) + self.logger.experiment.add_image(f"Train/Images_labels_{label_str}", grid, self.current_epoch) + except Exception as e: + print(f"Error logging train images: {e}") + + # ------------------------------------------------------------------ + # Predictions based on one backbone or two back bones + # ------------------------------------------------------------------ + if not self.is_difussion: + if self.num_inputs == 2: + peak_time = features["peak_time"] + peak_time = peak_time.to(self.device) + classification_pred, energy_pred, direction_pred = self.model( + imgs, peak_time + ) + else: + classification_pred, energy_pred, direction_pred = self.model(imgs) + + # ------------------------------------------------------------------ + # Particle type + # --------------------------------------- + if self.task == Task.type: + if self.is_difussion: + loss, accuracy, predicted, precision = self.compute_type_loss_diffusion(imgs,labels_class,training=True) + + else: + classification_pred_ = classification_pred[0] + feature_vector = classification_pred[1] + + loss, accuracy, predicted, precision = self.compute_type_loss( + classification_pred_, labels_class, test_val=False, training=True + ) + # Log batch loss and accuracy on the progress bar + self.log( + "train_acc", + accuracy * 100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + ) + self.log( + "train_precision", + precision*100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + ) + + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection: + + if self.is_difussion: + if self.num_inputs == 2: + peak_time = features["peak_time"] + peak_time = peak_time.to(self.device) + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff = self.compute_camera_direction_loss_diffusion(imgs, labels_direction,labels_energy_value, peak_time) + else: + + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff = self.compute_camera_direction_loss_diffusion(imgs, labels_direction,labels_energy_value) + else: + if len(direction_pred)==2: + direction_pred = direction_pred[0] + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff = self.compute_camera_direction_loss( direction_pred, labels_direction, training=True) + + self.loss_train_distance += loss_distance.item() + self.loss_train_dx_dy += loss_dx_dy.item() + self.loss_train_angular_error += loss_angular_diff.item() + + self.loss_val_separation = 0.0 + self.loss_val_dx_dy = 0.0 + self.loss_val_angular_error = 0.0 + + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy: + + if self.is_difussion: + if self.num_inputs == 1: + loss, *_ = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=True,x_2=None) + else: + peak_time = features["peak_time"] + loss, *_ = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=True,x_2=peak_time) + + else: + if len(energy_pred)==2: + energy_pred = energy_pred[0] + + loss, *_ = self.compute_energy_loss( + energy_pred, labels_energy_value, test_val=False, training=True + ) + # --------------------------------------- + + # L2 Regularization + l2_norm = sum(p.pow(2.0).sum() for p in self.parameters()) + loss = loss + self.l2_lambda * l2_norm + + # Display on Bar progress + self.log( + "train_loss", + loss.item(), + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + ) + + if not np.isnan(loss.item()): + + self.loss_train_sum += loss.item() + self.num_train_batches += 1 + + # torch.cuda.empty_cache() + return loss + # ---------------------------------------------------------------------------------------------------------- + def on_train_epoch_end(self): + + if self.trainer.world_size > 1: + + # dist.barrier() + + # Gather y Sync values with the GPUs. + total_loss_train= self.all_gather(self.loss_train_sum).sum().item() + total_batches_val = self.all_gather(torch.tensor(self.num_train_batches, device=self.device)).sum().item() + + else: + total_loss_train = self.loss_train_sum + total_batches_val = self.num_train_batches + + if total_batches_val==0: + + return 0 + + if self.trainer.is_global_zero: + + global_loss = total_loss_train / total_batches_val + + self.logger.experiment.add_scalars( + "loss/Global Training Loss", + { + "loss": global_loss, + }, + self.current_epoch, + ) + + self.logger.experiment.add_scalars( + "Learning rate", + { + "Learning rate": self.scheduler.get_last_lr()[0], + }, + self.current_epoch, + ) + # Optionally print the global loss for immediate feedback + print(f"Epoch {self.current_epoch}: Global Training Loss: {global_loss:.4f}") + + # --------------------------------------- + # Particle Type + # --------------------------------------- + if self.task == Task.type: + ii = 0 + # f1_score = self.f1_score_train.compute().detach().cpu().numpy()*100.0 + # self.f1_score_train.reset() + + # precision = self.precision_train.compute().detach().cpu().numpy()*100.0 + # self.precision_train.reset() + + epoch_accuracy = self.class_train_accuracy.compute().detach().cpu().item() * 100 + self.class_train_accuracy.reset() + + # # Compute the accuracy and reset the metric states after each epoch + if self.trainer.is_global_zero: + # self.log("train_acc_epoch", epoch_accuracy, on_step=False, prog_bar=True, sync_dist=True) + # # Log + # self.logger.experiment.add_scalars( + # "Metrics/Training", + # { + # "acc": epoch_accuracy, + # "f1":f1_score, + # "precision":precision, + # }, + # self.current_epoch, + # ) + print( + f"Epoch {self.current_epoch}: Global Training Accuracy: {epoch_accuracy:.4f}" + ) + filename_prefix = ( + f"Epoch_{self.current_epoch}_{self.task.name}_train_acc" + ) + if self.logger.log_dir: + self.save_checkpoint( + self.logger.log_dir, + epoch_accuracy, + filename_prefix=filename_prefix, + is_loss=False, + ) + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection and self.trainer.is_global_zero: + + filename_prefix = ( + f"Epoch_{self.current_epoch}_{self.task.name}_train_loss" + ) + self.save_checkpoint( + self.logger.log_dir, + global_loss, + filename_prefix=filename_prefix, + is_loss=True, + ) + # Log scalar values + self.logger.experiment.add_scalars( + "loss/ Loss Training", + { + "loss": global_loss, + "loss_distance": self.loss_train_distance + / self.num_train_batches, + "loss_alt_az": self.loss_train_dx_dy / self.num_train_batches, + "loss_angular_error": self.loss_train_angular_error + / self.num_train_batches, + }, + self.current_epoch, + ) + self.loss_train_distance = 0.0 + self.loss_train_dx_dy = 0.0 + self.loss_train_angular_error = 0.0 + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy and self.trainer.is_global_zero: + filename_prefix = ( + f"Epoch_{self.current_epoch}_{self.task.name}_train_loss" + ) + self.save_checkpoint( + self.logger.log_dir, + global_loss, + filename_prefix=filename_prefix, + is_loss=True, + ) + # --------------------------------------- + # Delete all the lists and set to 0 + # the values used to estimate the losses + # --------------------------------------- + self.reset_values() + print("END train epoch") + # Reset + self.loss_train_sum = 0 + self.num_train_batches = 0 + self.training_step_outputs = [] + # ---------------------------------------------------------------------------------------------------------- + @torch.no_grad() + def validation_step(self, batch, batch_idx, dataloader_idx=0): + loss=0 + self.model.eval() + + # ------------------------------------------------------------------ + # Read inputs (features) and labels + # ------------------------------------------------------------------ + features, labels, t = batch + + if len(features) > 0: + + imgs = features["image"] + + batch_size = imgs.shape[0] + if self.task == Task.type: + labels_class = labels["type"] + + labels_energy_value = labels["energy"] + hillas_intensity = features["hillas"]["hillas_intensity"] + + if self.task == Task.cameradirection: + labels_direction = labels["direction"] + + + # Log sample images to TensorBoard every 5 epochs + if batch_idx == 0 and self.current_epoch % 5 == 0 and self.trainer.is_global_zero: + import torchvision + try: + num_samples = min(4, imgs.shape[0]) + sample_imgs = imgs[:num_samples].to(self.device) + + label_str = "" + if self.task == Task.type: + label_str = str(labels_class[:num_samples].tolist()) + elif self.task == Task.energy: + label_str = str(labels_energy_value[:num_samples].tolist()) + elif self.task == Task.cameradirection: + label_str = str(labels_direction[:num_samples].tolist()) + + grid = torchvision.utils.make_grid(sample_imgs, normalize=True) + self.logger.experiment.add_image(f"Validation/Images_labels_{label_str}", grid, self.current_epoch) + except Exception as e: + print(f"Error logging validation images: {e}") + # ------------------------------------------------------------------ + # Predictions based on one backbone or two back bones + # ------------------------------------------------------------------ + if not self.is_difussion: + if self.num_inputs == 2: + peak_time = features["peak_time"] + classification_pred, energy_pred, direction_pred = self.model( + imgs, peak_time + ) + else: + classification_pred, energy_pred, direction_pred = self.model(imgs) + # ------------------------------------------------------------------ + # Compute Loss functions based on different tasks + # ------------------------------------------------------------------ + # Particle Type + # --------------------------------------- + if self.task == Task.type: + if self.is_difussion: + loss, accuracy, predicted, precision = self.compute_type_loss_diffusion(imgs,labels_class,training=False) + else: + classification_pred_ = classification_pred[0] + feature_vector = classification_pred[1] + loss, accuracy, predicted, precision = self.compute_type_loss( + classification_pred_, + labels_class, + test_val=False, + training=False, + ) + self.val_proton_gammanes.extend( + torch.softmax(classification_pred_, dim=1)[:, 1][labels_class == 0] + .float() + .cpu() + .detach() + .numpy() + .flatten() + .tolist() + ) + self.val_gamma_gammanes.extend( + torch.softmax(classification_pred_, dim=1)[:, 1][labels_class == 1] + .float() + .cpu() + .detach() + .numpy() + .flatten() + .tolist() + ) + # Log batch loss and accuracy on the progress bar + # if dataloader_idx == 0: + # if self.is_difussion: + # loss, accuracy, predicted, precision =self.compute_type_loss_diffusion( + # classification_pred_, + # labels_class, + # test_val=False, + # training=False, + # ) + # else: + # loss, accuracy, predicted, precision = self.compute_type_loss( + # classification_pred_, + # labels_class, + # test_val=False, + # training=False, + # ) + + self.log( + "val_acc", + accuracy * 100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + batch_size=batch_size, + ) + self.log( + "val_prec", + precision*100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + batch_size=batch_size, + ) + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection: + + if len(direction_pred)==2: + direction_pred = direction_pred[0] + + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_error= self.compute_camera_direction_loss( direction_pred, labels_direction, training=False) + + # ------------------------------------------------------------------------ + # Convert the offset to altitud and azimuth + # ------------------------------------------------------------------------ + + if dataloader_idx == 0: + self.loss_val_distance += loss_distance.item() + self.loss_val_dx_dy += loss_dx_dy.item() + self.loss_val_angular_error += loss_angular_diff.item() + self.val_angular_diff_list.extend(angular_error) + # ------------------------------------------------------- + + pred_dx = direction_pred[:, 0].float().cpu().detach().numpy() + pred_dy = direction_pred[:, 1].float().cpu().detach().numpy() + + cam_x = pred_dx + cam_y = pred_dy + + pred_alt, pred_az = self.val_dataset.cam_to_alt_az(labels["tel_ids"].cpu().detach().numpy(), labels["focal_length"].cpu().detach().numpy(), labels["pix_rotation"].cpu().detach().numpy(),labels["tel_az"].cpu().detach().numpy(),labels["tel_alt"].cpu().detach().numpy(), cam_x, cam_y) + + labels_dx_dy = labels_direction[:, 0:2] + + true_alt, true_az = self.val_dataset.cam_to_alt_az(labels["tel_ids"].cpu().detach().numpy(), labels["focal_length"].cpu().detach().numpy(), labels["pix_rotation"].cpu().detach().numpy(),labels["tel_az"].cpu().detach().numpy(),labels["tel_alt"].cpu().detach().numpy(), labels_dx_dy[:,0].float().cpu().detach().numpy(), labels_dx_dy[:,1].float().cpu().detach().numpy()) + + self.val_alt_pred_list.extend(np.radians(pred_alt)) + self.val_az_pred_list.extend(np.radians(pred_az)) + self.val_alt_label_list.extend(np.radians(true_alt)) + self.val_az_label_list.extend(np.radians(true_az)) + + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy: + if self.is_difussion: + if self.num_inputs == 1: + loss, energy_diff, energy_pred_tev = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=False,x_2=None) + else: + peak_time = features["peak_time"] + loss, energy_diff, energy_pred_tev = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=False,x_2=peak_time) + else: + if len(energy_pred)==2: + energy_pred = energy_pred[0] + + loss, energy_diff = self.compute_energy_loss( + energy_pred, labels_energy_value/1.0, test_val=False, training=False + ) + energy_pred_tev = torch.pow(10, energy_pred*1.0) + + if dataloader_idx == 0: + + self.val_energy_diff_list.extend(energy_diff) + self.val_energy_pred_list.extend( + energy_pred_tev[:, 0].float().cpu().detach().numpy() + ) + + # --------------------------------------- + # Collect the True Energy and Hillas Intensity + # --------------------------------------- + + energy_label_tev = torch.pow(10, labels_energy_value) + + if dataloader_idx == 0: + self.val_energy_label_list.extend( + energy_label_tev[:, 0].float().cpu().detach().numpy().flatten().tolist() + ) + self.val_hillas_intensity_list.extend( + hillas_intensity[:, 0].float().cpu().detach().numpy().flatten().tolist() + ) + + # --------------------------------------- + # Log validation loss + # --------------------------------------- + + if dataloader_idx == 0: + self.loss_val_sum += loss.item() + self.num_val_batches += 1 + + + loss_key = "val_loss" if dataloader_idx == 0 else "test_loss" + if loss is not None: + self.log(loss_key, loss,on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size) + + # Release cuda memory + # torch.cuda.empty_cache() + return loss + # ---------------------------------------------------------------------------------------------------------- + @torch.no_grad() + def on_validation_epoch_end(self): + + if self.trainer.world_size > 1: + dist.barrier() + # Gather y Sync values with the GPUs. + total_loss_val = self.all_gather(self.loss_val_sum).sum().item() + total_batches_val = self.all_gather(torch.tensor(self.num_val_batches, device=self.device)).sum().item() + + else: + total_loss_val = self.loss_val_sum + total_batches_val = self.num_val_batches + + + if self.trainer.is_global_zero: + # Calcular la pérdida promedio global + global_loss_val = total_loss_val / max(1, total_batches_val) + + self.logger.experiment.add_scalars( + "loss/Global Validation Loss", + { + "loss": global_loss_val, + }, + self.current_epoch, + ) + + # Print the global loss for immediate feedback + print(f"Epoch {self.current_epoch}: Global Validation Loss: {global_loss_val:.4f}") + # --------------------------------------- + # Particle Type + # --------------------------------------- + if self.task == Task.type: + conf_matrix = self.confusion_matrix.compute().detach().cpu().numpy() + f1_score_val = self.f1_score_val.compute().detach().cpu().numpy()*100.0 + precision_val = self.precision_val.compute().detach().cpu().numpy()*100.0 + + # Compute the accuracy and reset the metric states after each epoch + epoch_accuracy_val = self.class_val_accuracy.compute().item() * 100 + + self.class_val_accuracy.reset() + self.confusion_matrix.reset() + self.f1_score_val.reset() + self.precision_val.reset() + + # --------------------------------------- + # Create Confusion Matrix + # --------------------------------------- + if self.trainer.is_global_zero: + + # Log + self.logger.experiment.add_scalars( + "Metrics/Validation", + { + "acc": epoch_accuracy_val, + "f1": f1_score_val, + "precision":precision_val, + }, + self.current_epoch, + ) + + print( + f"Epoch {self.current_epoch}: Global Validation Accuracy: {epoch_accuracy_val:.4f}" + ) + print( + f"Epoch {self.current_epoch}: Global Validation F1 Score: {f1_score_val:.4f}" + ) + print( + f"Epoch {self.current_epoch}: Global Validation Precision: {precision_val:.4f}" + ) + filename_prefix = "confusion_matrix_val" + cm_file_name = f"{filename_prefix}_{self.current_epoch}_{epoch_accuracy_val:.4f}_Validation" + + if self.logger.log_dir: + plot_confusion_matrix( + conf_matrix, + self.class_names, + cm_file_name, + self.logger.log_dir, + ) + # Compute class-wise accuracies + class_accuracies = (conf_matrix.diagonal() / conf_matrix.sum(axis=1))*100 + + self.logger.experiment.add_scalars( + "confusion_matrix", + { + "val_acc_gamma": class_accuracies[0], + "val_acc_proton": class_accuracies[1], + "val_global_acc": epoch_accuracy_val, + }, + self.current_epoch, + ) + # Add the protonness and gammaness histogramns + self.val_proton_gammanes = np.array(self.val_proton_gammanes) + self.val_gamma_gammanes = np.array(self.val_gamma_gammanes) + if self.trainer.is_global_zero: + fig, ax = plt.subplots(figsize=(8, 6)) + ax.hist( + self.val_proton_gammanes, + bins=50, + alpha=0.7, + label="Proton", + color="blue", + density=True, + ) + ax.hist( + self.val_gamma_gammanes, + bins=50, + alpha=0.7, + label="Gammas", + color="orange", + density=True, + ) + ax.set_xlabel("Score") + ax.set_ylabel("Density") + ax.set_title("Protons and Gammas Gammanes Distribution - Validation") + ax.legend() + save_path = os.path.join( + self.logger.log_dir, + f"proton_gammaness_validation_{self.current_epoch}_{global_loss_val}.png" + ) + fig.savefig(save_path, format="png") # <--- Usa fig.savefig, no plt.savefig + + # 2. Luego envíala al logger + self.logger.experiment.add_figure( + "Protons and Gammas Gammanes Distribution - Validation", + fig, + self.current_epoch, + ) + + # 3. Finalmente cierra la figura + plt.close(fig) + + # Clear lists + self.val_proton_gammanes = [] + self.val_gamma_gammanes = [] + # return 0 + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection and self.trainer.is_global_zero: + + self.print_direction_error(self.val_angular_diff_list, "Validation") + + plt.close("all") + fig_direction_error = plot_direction_resolution_error( + self.val_alt_pred_list, + self.val_az_pred_list, + self.val_alt_label_list, + self.val_az_label_list, + self.val_energy_label_list, + self.val_hillas_intensity_list, + ) + self.logger.experiment.add_figure( + "Direction Resolution Error/Validation", + fig_direction_error, + self.current_epoch, + ) # Log the plot + + # Save the figure + fig_direction_error.savefig( + os.path.join( + self.logger.log_dir, + "angular_resolution_validation_" + + str(self.current_epoch) + + "_" + + str(global_loss_val) + + ".png", + ), + format="png", + ) + plt.close(fig_direction_error) # Close the figure to release memory + plt.close("all") + + # Log scalar values + self.logger.experiment.add_scalars( + "loss/Loss Validation", + { + "loss": global_loss_val, + "loss_distance": self.loss_val_distance / self.num_val_batches, + "loss_dx_dy": self.loss_val_dx_dy / self.num_val_batches, + "loss_angular_error": self.loss_val_angular_error + / self.num_val_batches, + }, + self.current_epoch, + ) + + + + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy and self.trainer.is_global_zero: + + self.print_energy_error(self.val_energy_diff_list, "Validation") + plt.close("all") + fig_energy_error = plot_energy_resolution_error( + self.val_energy_pred_list, + self.val_energy_label_list, + self.val_hillas_intensity_list, + ) + self.logger.experiment.add_figure( + "Energy Resolution Error/Validation", + fig_energy_error, + self.current_epoch, + ) # Log the plot + fig_energy_error.savefig( + os.path.join( + self.logger.log_dir, + "error_resolution_validation_" + + str(self.current_epoch) + + "_" + + str(global_loss_val) + + ".png", + ), + format="png", + ) + + plt.close(fig_energy_error) # Close the figure to release memory + plt.close("all") + + # Add histograms for energy labels and predictions + fig_hist, ax_hist = plt.subplots(figsize=(8, 6)) + + import numpy as np + # Use percentiles of true labels to avoid extreme predictions ruining the x-axis range + min_val = np.percentile(self.val_energy_label_list, 1) + max_val = np.percentile(self.val_energy_label_list, 99) + + # Energy typically spans multiple orders of magnitude, so log bins are more appropriate + min_val = max(1e-3, min_val) + if min_val >= max_val: + max_val = min_val * 10 + shared_bins = np.logspace(np.log10(min_val), np.log10(max_val), 50) + + ax_hist.hist( + self.val_energy_label_list, + bins=shared_bins, + alpha=0.5, + label="Labels", + color="blue", + ) + ax_hist.hist( + self.val_energy_pred_list, + bins=shared_bins, + alpha=0.5, + label="Predictions", + color="orange", + ) + ax_hist.set_xscale("log") + ax_hist.set_xlabel("Energy (TeV)") + ax_hist.set_ylabel("Counts") + ax_hist.set_title("Energy Distribution - Labels vs Predictions (Validation)") + ax_hist.legend() + + self.logger.experiment.add_figure( + "Energy Distribution - Labels vs Predictions/Validation", + fig_hist, + self.current_epoch, + ) + + save_path = os.path.join( + self.logger.log_dir, + f"energy_distribution_validation_{self.current_epoch}_{global_loss_val}.png" + ) + fig_hist.savefig(save_path, format="png") + plt.close(fig_hist) + plt.close("all") + + # ---------------------------------------------------------------------------------------------------------- + def reset_values(self): + # --------------------------------------- + # Reset + # --------------------------------------- + self.loss_val_sum = 0 + self.num_val_batches = 0 + + # Reset Energy and hillas intensity + # -------------------------------------------------- + # Validation + # -------------------------------------------------- + self.val_energy_label_list.clear() + self.val_hillas_intensity_list.clear() + + # -------------------------------------------------- + # Reset Direction + # -------------------------------------------------- + # Validation + # -------------------------------------------------- + self.val_angular_diff_list.clear() + + + self.val_alt_pred_list.clear() + self.val_az_pred_list.clear() + self.val_alt_label_list.clear() + self.val_az_label_list.clear() + + self.loss_val_distance = 0.0 + self.loss_val_dx_dy = 0.0 + self.loss_val_angular_error = 0.0 + + # -------------------------------------------------- + # Reset Energy + # -------------------------------------------------- + # Validation + # -------------------------------------------------- + self.val_energy_diff_list.clear() + self.val_energy_pred_list.clear() + + + # ---------------------------------------------------------------------------------------------------------- + def print_direction_error(self, angular_diff_list, type_val: str): + # Count the angular error in ranges [20º-0.1º] + error_20 = len([num for num in angular_diff_list if num < 20]) + error_10 = len([num for num in angular_diff_list if num < 10]) + error_5 = len([num for num in angular_diff_list if num < 5]) + error_2 = len([num for num in angular_diff_list if num < 2]) + error_1 = len([num for num in angular_diff_list if num < 1]) + error_0_5 = len([num for num in angular_diff_list if num < 0.5]) + error_0_25 = len([num for num in angular_diff_list if num < 0.25]) + error_0_1 = len([num for num in angular_diff_list if num < 0.1]) + # Log + self.logger.experiment.add_scalars( + "Direction Error/" + type_val, + { + "0: error: 20": error_20, + "1: error: 10": error_10, + "2: error: 5": error_5, + "3: error: 2": error_2, + "4: error: 1": error_1, + "5: error: 0.5": error_0_5, + "6: error: 0.25": error_0_25, + "7: error: 0.1": error_0_1, + }, + self.current_epoch, + ) + # Print + print(type_val + " Direction Error < 20º: ", error_20) + print(type_val + " Direction Error < 10º: ", error_10) + print(type_val + " Direction Error < 5º: ", error_5) + print(type_val + " Direction Error < 2º: ", error_2) + print(type_val + " Direction Error < 1º: ", error_1) + print(type_val + " Direction Error < 0.50º: ", error_0_5) + print(type_val + " Direction Error < 0.25º: ", error_0_25) + print(type_val + " Direction Error < 0.10º: ", error_0_1) + # ---------------------------------------------------------------------------------------------------------- + def print_energy_error(self, energy_diff_list, type_val: str): + + + error_30 = len([num for num in energy_diff_list if num < 30]) + error_20 = len([num for num in energy_diff_list if num < 20]) + error_10 = len([num for num in energy_diff_list if num < 10]) + error_5 = len([num for num in energy_diff_list if num < 5]) + error_1 = len([num for num in energy_diff_list if num < 1]) + error_0_5 = len([num for num in energy_diff_list if num < 0.5]) + error_0_25 = len([num for num in energy_diff_list if num < 0.25]) + error_0_1 = len([num for num in energy_diff_list if num < 0.1]) + error_0_01 = len([num for num in energy_diff_list if num < 0.01]) + # Log + self.logger.experiment.add_scalars( + "Energy Error/" + type_val, + { + "0: error: 30": error_30, + "1: error: 20": error_20, + "2: error: 10": error_10, + "3: error: 5": error_5, + "4: error: 1": error_1, + "5: error: 0.5": error_0_5, + "6: error: 0.25": error_0_25, + "7: error: 0.1": error_0_1, + "8: error: 0.01": error_0_01, + }, + self.current_epoch, + ) + # Print + print(f"Total: {len(energy_diff_list)}") + print(type_val + " Energy Error < 30:", error_30) + print(type_val + " Energy Error < 20:", error_20) + print(type_val + " Energy Error < 10:", error_10) + print(type_val + " Energy Error < 5:", error_5) + print(type_val + " Energy Error < 1:", error_1) + print(type_val + " Energy Error < 0.5:", error_0_5) + print(type_val + " Energy Error < 0.1:", error_0_1) + print(type_val + " Energy Error < 0.01:", error_0_01) + # ---------------------------------------------------------------------------------------------------------- + def create_confusion_matrix( + self, epoch, accuracy, all_val_preds, all_val_labels, val_type: str + ): + all_labels = all_val_preds.numpy() + all_preds = all_val_labels.numpy() + filename_prefix = "confusion_matrix_val" + cm_file_name = f"{filename_prefix}_{epoch}_{accuracy:.4f}_{val_type}" + cm = confusion_matrix(all_labels, all_preds) + + cm_norm = cm.astype("float") / cm.sum(axis=1)[:, np.newaxis] + + # Remove NaNs + cm_norm = np.nan_to_num(cm_norm, nan=0.0) + accuracies = cm_norm.diagonal() * 100.0 + if self.logger.log_dir: + plot_confusion_matrix( + cm, + self.class_names, + accuracies, + cm_file_name, + self.logger.log_dir, # self.save_folder + ) + + return accuracies + # ---------------------------------------------------------------------------------------------------------- + def configure_optimizers(self): + + if self.optimizer_type == "sgd": + print("Using SGD...") + # Optimizer SGD + self.optimizer = optim.SGD( + self.model.parameters(), + lr=self.learning_rate, + momentum=self.momentum, + weight_decay=self.weight_decay, + nesterov=True, + ) + elif self.optimizer_type == "adam": + print("Using Adam...") + # Optimizer Adam + self.optimizer = optim.Adam( + self.model.parameters(), + lr=self.learning_rate, + eps=self.adam_epsilon, + weight_decay=self.weight_decay, + ) + elif self.optimizer_type == "adamw": + print("Using AdamW...") + # Optimizer Adam + self.optimizer = optim.AdamW( + self.model.parameters(), + lr=self.learning_rate, + eps=self.adam_epsilon, + weight_decay=self.weight_decay, + ) + else: + raise ValueError( + f"Unsupported optimizer type: {self.optimizer_type}. Supported types are 'sgd' and 'adam'." + ) + + # Cosine phase + # Note: The schedule is updated after execute one epoch. -> self.epochs* self.batches + lf = one_cycle( + 1, self.lrf, self.num_epochs * self.steps_epoch + ) # cosine 1->hyp['lrf'] + # self.scheduler = lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lf) + + # TODO: Add the option to choose the kind of scheduler + # self.scheduler = torch.optim.lr_scheduler.OneCycleLR(self.optimizer,max_lr=self.lrf,total_steps=self.num_epochs * self.steps_epoch) + self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + self.optimizer, T_max=self.num_epochs + ) + self.scheduler.last_epoch = self.start_epoch + # Prepare scheduler dictionary as expected by PyTorch Lightning + lr_dict = { + "scheduler": self.scheduler, + "interval": "epoch", # Use for each epoch, dont use for each step + "frequency": 1, + "name": "OneCycleLR", + "monitor": "val_loss", + } + # Return the optimizer and learning rate scheduler + return { + "optimizer": self.optimizer, + "lr_scheduler": lr_dict, + } + # ---------------------------------------------------------------------------------------------------------- diff --git a/ctlearn/tools/pytorch/__init__.py b/ctlearn/tools/pytorch/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/tools/pytorch/config/default_config_file.yml b/ctlearn/tools/pytorch/config/default_config_file.yml new file mode 100644 index 00000000..abf55871 --- /dev/null +++ b/ctlearn/tools/pytorch/config/default_config_file.yml @@ -0,0 +1,152 @@ +data: + train_gamma_proton: ./data/gamma_proton_train_remix.dl1.pickle + validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle + + train_gamma: ./data/gamma_955000_train.pickle + validation_gamma: ./data/gamma_106141_validation.pickle + + test_gamma: ./data/gamma_1805522_test_gamma.pickle + test_proton: ./data/proton_130811_test_proton.pickle + test_electron: None + test_validation_gamma: ./data/gamma_180552_test_val_gamma.pickle + + test_validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle + + observation: ./run_2931.dl1.pickle + # Important: This is only for testing purpose. Set always to 0 + # when you are training, validating or estimating the dl2 files + training_reduce_factor: 0 #64 #4 + validation_reduce_factor: 0 #16 #8 + validation_test_reduce_factor: 0 #16 #8 + + # Check points + type_checkpoint: ./run/run_type_training_14/exp_14_type_train/version_0/Epoch_6_type_train_acc_80.9682309627532959.pth + energy_checkpoint: /home/cpozogonzalez/ctlearn/run/run_energy_training_14/exp_14_energy_train/version_1/Epoch_13_energy_train_loss_30.2185532478501244.pth + direction_checkpoint: /lhome/ext/ucm147/ucm1477/data/check_points/v_5/Epoch_23_cameradirection_train_loss_7296.6534562211982120.pth + +run_details: + + mode: "observation" # The option are: "train", "results", "observation" and "validate" + task: "direction" # The option are: "all", "energy" "type" and "direction" + test_type: "gamma" # The option are: "gamma" "proton" or "electron" + experiment_number: 14 # The experiment number. The experiment folder is saved into the "run" folder. + + +cut-off: + + leakage_intensity: 0.2 # bigger to this value, the event is removed + intensity: 50 # below to this value, the event is removed + +model: + + model_type: + model_name: "DoubleBBEfficientNet" + parameters: + model_variant: "efficientnet-b3" + task: 'type' + num_outputs: 2 + device_str: "cuda" + energy_bins: None + + model_energy: + model_name: "ThinResNet" + parameters: + task: 'energy' + num_inputs: 1 + num_outputs: 1 + num_blocks: [3, 4, 6, 3] #[2, 3, 3, 3] + dropout: 0.1 + use_bn: False + + model_direction: + model_name: "ThinResNet_DBB" + parameters: + task: 'direction' + num_inputs: 1 + num_outputs: 3 + num_blocks: [3, 4, 6, 3] + dropout: 0.1 + use_bn: False + + # model_direction: + # model_name: "DBBNoPropDTReg" + # parameters: + # task: 'direction' + # num_outputs: 3 + # embedding_dim: 512 + # T: 3 + # eta: 0.1 #0.1 + +# Hyper-parameters +hyp: + + epochs: 30 + batches: 128 #128 #64 + dynamic_batches: True + optimizer: Adamw + momentum: 0.957 #Yolo 0.937 # Efficient-b3 0.757 + weight_decay: 0.0005 #0.004676 #0.0001 #0.00002 Efficient-b3 0.0005 + learning_rate: 1e-4 #1e-5 #Efficient-b3 1e-5 + lrf: 0.1 + start_epoch: 0 + steps_epoch: 100 # Computed online. Must be removed + l2_lambda: 1e-7 #1e-5 #1e-5 # L2 regularization (Set to 0.0 to skip the L2 Regularization) + adam_epsilon: 1.0e-08 #7.511309034256153e-05 #1.0e-08 + gradient_clip_val: 3.0 # Avoid gradient explosion + + save_k: 200 # Save as maximum k checkpoints. + +augmentation: + # probabilities for augmentation range = [0, 1.0] + # prob = 0.0 -> Always apply the augmentation + # prob >= 1.0 -> Never apply the augmentation, i.e., Set bigger than 1.0 ( ex: 2.0) if you want disable it. + # Note: mask augmentation is always on even with flag use_augmentation = True + # To disable it, just set to 2.5 for example. + + use_augmentation: True # This apply only on training mode. + aug_prob: 0.5 # Probability of use Augmentation + rot_prob: 0.5 # Rotation probability + trans_prob: 0.5 # Translation probability + flip_hor_prob: 0.5 # Horizontal Flip probability + flip_ver_prob: 0.5 # Vertical Flip probability + mask_prob: 0.5 # Apply mask probability + mask_dvr_prob: 0.5 # Apply dvr mask probability + noise_prob: 0.5 # No implemented yet. + max_rot: 5 # Maximum rotation in augmentation + max_trans: 10 # Maximum translation in augmentation + +normalization: + + # Normalization: Im' = (Im-mu)/sigma + use_clean: True # Use the image with the applied mask (True), IOC the mask is not applied (False) + use_clean_dvr: False + type_mu: 0.0 + type_sigma: 1000.0 + + dir_mu: 0.0 + dir_sigma: 1000.0 + + energy_mu: 0.0 + energy_sigma: 1000.0 + +dataset: + num_workers: 1 # + pin_memory: True + persistent_workers: True # + +# Hardware Architecture and precision +arch: + # device: 'mps' # Apple Mx + device: 'cuda' + precision_type: "32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" + precision_energy: "32-true" #"32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" + precision_direction: "32-true" # "bf16-mixed" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" + # (bf16 for GPU with Ampere or higher, it is better that 16 because is numerical more stability) + # devices: [0,1] # [0,1] For multiple GPUs + devices: [0,1] + # Note: Check the documentation for more information. + strategy: 'deepspeed_stage_2' # Options: auto, dpp, dpp_swap, fsdp, deepspeed, horovod, bagua, deepspeed_stage_2, deepspeed_stage_3, colossalai, hivemind, etc... + +Notes: + Note_1: Training with augmentation dvr using 1-3 dilatations + Note_2: Trainining b3 applying always the mask \ No newline at end of file diff --git a/ctlearn/tools/pytorch/train_pytorch_model.py b/ctlearn/tools/pytorch/train_pytorch_model.py new file mode 100644 index 00000000..6156afa9 --- /dev/null +++ b/ctlearn/tools/pytorch/train_pytorch_model.py @@ -0,0 +1,710 @@ +from ctapipe.core.traits import Path, Bool, Unicode +from torch.utils.data import DataLoader +from pytorch_ctlearn.ctlearn.ctlearn.tools.pytorch.CTLearnPL import CTLearnTrainer, CTLearnPL +import os +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Suppress TensorFlow logging + +try: + import torch + +except ImportError: + raise ImportError("pytorch is not installed in your environment!") + +try: + from pytorch_lightning.loggers import TensorBoardLogger +except ImportError: + raise ImportError("pytorch_lightning is not installed in your environment!") + +from ctlearn.tools.train.base_train_model import TrainCTLearnModel +from ctlearn.core.ctlearn_enum import Task, Mode +from .utils import ( + str_list_to_enum_list, + sanity_check, + read_configuration, + create_experiment_folder, + expected_structure, +) + +from ctlearn.core.pytorch.net_utils import create_model, ModelHelper +from ctlearn.core.data_loader.loader import DLDataLoader +from pytorch_lightning.callbacks import Callback +import os +import numpy as np +import json + + +class GPUStatsLogger(Callback): + def on_train_epoch_end(self, trainer, pl_module): + mem_allocated = torch.cuda.memory_allocated() + mem_reserved = torch.cuda.memory_reserved() + + trainer.logger.experiment.add_scalar( + "gpu_mem_allocated", mem_allocated, global_step=trainer.current_epoch + ) + trainer.logger.experiment.add_scalar( + "gpu_mem_reserved", mem_reserved, global_step=trainer.current_epoch + ) + +# from ctlearn.tools.train_model import +class TrainPyTorchModel(TrainCTLearnModel): + """ + Tool to train a ``~ctlearn.core.model.CTLearnModel`` on R1/DL1a data using PyTorch. + + The tool sets up the PyTorch model using ... The PyTorch model is trained + on the input data (R1 calibrated waveforms or DL1a images) and saved in the output directory. + """ + + name = "ctlearn-train-pytorch-model" + description = __doc__ + + examples = """ + To train a CTLearn PyTorch model for the classification of the primary particle type: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --background /path/to/your/protons_dl1_dir/ \\ + --pattern-background "proton_*_run1.dl1.h5" \\ + --pattern-background "proton_*_run10.dl1.h5" \\ + --output /path/to/your/type/ \\ + --reco type \\ + + To train a CTLearn PyTorch model for the regression of the primary particle energy: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/energy/ \\ + --reco energy \\ + + To train a CTLearn PyTorch model for the regression of the primary particle + arrival direction based on the offsets in camera coordinates: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/direction/ \\ + --reco cameradirection \\ + + To train a CTLearn PyTorch model for the regression of the primary particle + arrival direction based on the offsets in sky coordinates: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/direction/ \\ + --reco skydirection \\ + """ + + config_file = Path( + exits=True, + default_value=None, + allow_none=True, + directory_ok=True, + file_ok=True, + help="Configuration file.", + ).tag(config=True) + + disable_progress_bar = Bool( + default_value=False, + help="Disable PyTorch Lightning progress bar.", + ).tag(config=True) + + model_name = Unicode( + default_value=None, + allow_none=True, + help="Model name to override the default model for the reco task.", + ).tag(config=True) + + aliases = { + **TrainCTLearnModel.aliases, + "config_file": "TrainPyTorchModel.config_file", + "disable_progress_bar": "TrainPyTorchModel.disable_progress_bar", + "model-name": "TrainPyTorchModel.model_name", + } + + def __init__(self, **kwargs): + + # Setup GPU + os.environ["NCCL_P2P_DISABLE"] = "1" + os.environ["NCCL_IB_DISABLE"] = "1" + os.environ["NCCL_DEBUG"] = "WARN" + os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1" + os.environ["NCCL_DEBUG"] = "INFO" + torch.set_float32_matmul_precision('medium') + + super().__init__(**kwargs) + + + def setup(self): + + super().setup() + + # Create tasks Enum List + self.tasks = str_list_to_enum_list(self.reco_tasks) + + for task_ in self.tasks: + print("Task:", task_.name) + + if self.config_file is not None: + self.log.info("Loading configuration from legacy PyTorch config file: %s", self.config_file) + legacy_params = read_configuration(self.config_file) + sanity_check(legacy_params, expected_structure) + + def get_conf_val(key1, key2, trait_name, default_val): + in_config = False + for cls_name in ["TrainPyTorchModel", "TrainCTLearnModel"]: + if cls_name in self.config and trait_name in self.config[cls_name]: + in_config = True + break + if not in_config: + return legacy_params.get(key1, {}).get(key2, default_val) + return getattr(self, trait_name) + + self.type_checkpoint = get_conf_val("data", "type_checkpoint", "type_checkpoint", self.type_checkpoint) + self.energy_checkpoint = get_conf_val("data", "energy_checkpoint", "energy_checkpoint", self.energy_checkpoint) + self.direction_checkpoint = get_conf_val("data", "direction_checkpoint", "direction_checkpoint", self.direction_checkpoint) + self.load_onnx_model = get_conf_val("data", "load_onnx_model", "load_onnx_model", self.load_onnx_model) + + self.experiment_number = get_conf_val("run_details", "experiment_number", "experiment_number", self.experiment_number) + self.save_onnx = get_conf_val("run_details", "save_onnx", "save_onnx", self.save_onnx) + self.save_k_checkpoints = get_conf_val("hyp", "save_k", "save_k_checkpoints", self.save_k_checkpoints) + self.device_str = get_conf_val("arch", "device", "device", self.device) + self.batch_size = get_conf_val("hyp", "batches", "batch_size", self.batch_size) + self.pin_memory = get_conf_val("dataset", "pin_memory", "pin_memory", self.pin_memory) + self.num_workers = get_conf_val("dataset", "num_workers", "num_workers", self.num_workers) + self.persistent_workers = get_conf_val("dataset", "persistent_workers", "persistent_workers", self.persistent_workers) + self.devices = get_conf_val("arch", "devices", "devices", self.devices) + self.strategy = get_conf_val("arch", "strategy", "strategy", self.strategy) + + # Augmentations + self.use_augmentation = get_conf_val("augmentation", "use_augmentation", "use_augmentation", self.use_augmentation) + self.aug_prob = get_conf_val("augmentation", "aug_prob", "aug_prob", self.aug_prob) + self.rot_prob = get_conf_val("augmentation", "rot_prob", "rot_prob", self.rot_prob) + self.trans_prob = get_conf_val("augmentation", "trans_prob", "trans_prob", self.trans_prob) + self.flip_hor_prob = get_conf_val("augmentation", "flip_hor_prob", "flip_hor_prob", self.flip_hor_prob) + self.flip_ver_prob = get_conf_val("augmentation", "flip_ver_prob", "flip_ver_prob", self.flip_ver_prob) + self.mask_prob = get_conf_val("augmentation", "mask_prob", "mask_prob", self.mask_prob) + self.mask_dvr_prob = get_conf_val("augmentation", "mask_dvr_prob", "mask_dvr_prob", self.mask_dvr_prob) + self.noise_prob = get_conf_val("augmentation", "noise_prob", "noise_prob", self.noise_prob) + self.max_rot = get_conf_val("augmentation", "max_rot", "max_rot", self.max_rot) + self.max_trans = get_conf_val("augmentation", "max_trans", "max_trans", self.max_trans) + + # Normalizations + self.use_clean = get_conf_val("normalization", "use_clean", "use_clean", self.use_clean) + self.use_clean_dvr = get_conf_val("normalization", "use_clean_dvr", "use_clean_dvr", self.use_clean_dvr) + self.type_mu = get_conf_val("normalization", "type_mu", "type_mu", self.type_mu) + self.type_sigma = get_conf_val("normalization", "type_sigma", "type_sigma", self.type_sigma) + self.dir_mu = get_conf_val("normalization", "dir_mu", "dir_mu", self.dir_mu) + self.dir_sigma = get_conf_val("normalization", "dir_sigma", "dir_sigma", self.dir_sigma) + self.energy_mu = get_conf_val("normalization", "energy_mu", "energy_mu", self.energy_mu) + self.energy_sigma = get_conf_val("normalization", "energy_sigma", "energy_sigma", self.energy_sigma) + + # Cut-offs + self.leakage_intensity_cutoff = get_conf_val("cut-off", "leakage_intensity", "leakage_intensity_cutoff", self.leakage_intensity_cutoff) + self.intensity_cutoff = get_conf_val("cut-off", "intensity", "intensity_cutoff", self.intensity_cutoff) + + self.pytorch_model_configs = legacy_params.get("model", {}) + self.hyp_configs = legacy_params.get("hyp", {}) + else: + self.log.info("No legacy config file provided. Using standard Traitlets configuration.") + self.device_str = self.device + self.pytorch_model_configs = { + "model_type": { + "model_name": "DoubleBBEfficientNet", + "parameters": { + "model_variant": "efficientnet-b3", + "task": "type", + "num_outputs": 2, + "device_str": self.device_str, + "energy_bins": None, + } + }, + "model_energy": { + "model_name": "ThinResNet", + "parameters": { + "task": "energy", + "num_inputs": 1, + "num_outputs": 1, + "num_blocks": [3, 4, 6, 3], + "dropout": 0.1, + "use_bn": False, + } + }, + "model_direction": { + "model_name": "ThinResNet_DBB", + "parameters": { + "task": "direction", + "num_inputs": 1, + "num_outputs": 3, + "num_blocks": [3, 4, 6, 3], + "dropout": 0.1, + "use_bn": False, + } + } + } + self.hyp_configs = { + "epochs": self.n_epochs, + "batches": self.batch_size, + "dynamic_batches": True, + "optimizer": self.optimizer.get("name", "Adam"), + "momentum": self.optimizer_momentum, + "weight_decay": self.optimizer_weight_decay, + "learning_rate": self.optimizer.get("base_learning_rate", 0.0001), + "lrf": self.lrf, + "start_epoch": 0, + "steps_epoch": 100, + "l2_lambda": self.l2_lambda, + "adam_epsilon": self.optimizer.get("adam_epsilon", 1.0e-8), + "gradient_clip_val": self.gradient_clip_val, + "save_k": self.save_k_checkpoints, + } + + # Override model name if passed through command line + if self.model_name is not None: + for task in self.tasks: + task_key = f"model_{task.name}" + if task_key not in self.pytorch_model_configs: + self.pytorch_model_configs[task_key] = {"parameters": {}} + self.pytorch_model_configs[task_key]["model_name"] = self.model_name + + # Clear out parameters to avoid passing the old model's params to the new one + if "parameters" in self.pytorch_model_configs[task_key]: + old_params = self.pytorch_model_configs[task_key]["parameters"] + # Keep basic essential params if they exist + new_params = {} + for key in ["task", "num_inputs", "num_outputs", "device_str"]: + if key in old_params: + new_params[key] = old_params[key] + self.pytorch_model_configs[task_key]["parameters"] = new_params + + self.save_k = self.save_k_checkpoints + + self.parameters = { + "data": { + "train_gamma_proton": None, + "validation_gamma_proton": None, + "train_gamma": None, + "validation_gamma": None, + "test_gamma": None, + "test_proton": None, + "test_electron": None, + "test_validation_gamma": None, + "test_validation_gamma_proton": None, + "type_checkpoint": self.type_checkpoint, + "energy_checkpoint": self.energy_checkpoint, + "direction_checkpoint": self.direction_checkpoint, + "load_onnx_model": self.load_onnx_model, + }, + "run_details": { + "mode": "train", + "task": self.reco_tasks[0] if self.reco_tasks else "all", + "test_type": "gamma", + "experiment_number": self.experiment_number, + }, + "cut-off": { + "leakage_intensity": self.leakage_intensity_cutoff, + "intensity": self.intensity_cutoff, + }, + "model": self.pytorch_model_configs, + "hyp": self.hyp_configs, + "augmentation": { + "use_augmentation": self.use_augmentation, + "aug_prob": self.aug_prob, + "rot_prob": self.rot_prob, + "trans_prob": self.trans_prob, + "flip_hor_prob": self.flip_hor_prob, + "flip_ver_prob": self.flip_ver_prob, + "mask_prob": self.mask_prob, + "mask_dvr_prob": self.mask_dvr_prob, + "noise_prob": self.noise_prob, + "max_rot": self.max_rot, + "max_trans": self.max_trans, + }, + "normalization": { + "apply_log_scaling": self.apply_log_scaling, + "use_clean": self.use_clean, + "use_clean_dvr": self.use_clean_dvr, + "type_mu": self.type_mu, + "type_sigma": self.type_sigma, + "dir_mu": self.dir_mu, + "dir_sigma": self.dir_sigma, + "energy_mu": self.energy_mu, + "energy_sigma": self.energy_sigma, + }, + "dataset": { + "num_workers": self.num_workers, + "pin_memory": self.pin_memory, + "persistent_workers": self.persistent_workers, + }, + "arch": { + "device": self.device_str, + "precision_type": self.precision_type, + "precision_energy": self.precision_energy, + "precision_direction": self.precision_direction, + "devices": self.devices, + "strategy": self.strategy, + } + } + + print(f"Using Devices: {self.devices}") + + # Set up the data loaders for training and validation + indices = list(range(self.dl1dh_reader._get_n_events())) + # Shuffle the indices before the training/validation split + np.random.seed(self.random_seed) + np.random.shuffle(indices) + n_validation_examples = int( + self.validation_split * self.dl1dh_reader._get_n_events() + ) + training_indices = indices[n_validation_examples:] + validation_indices = indices[:n_validation_examples] + + if not ("class_weight" in self.parameters): + self.parameters['class_weight'] = self.dl1dh_reader.class_weight + self.log.info(f"Class weights not provided. Using class weights from data reader: {self.parameters['class_weight']}") + elif len(self.parameters['class_weight']) != len(self.dl1dh_reader.class_names): + raise ValueError(f"Number of class weights provided ({len(self.parameters['class_weight'])}) does not match number of classes in data ({len(self.dl1dh_reader.class_names)}).") + else: + self.log.info(f"Using class weights from configuration file: {self.parameters['class_weight']}") + + print("BASE TRAIN FRAMEWORK", self.framework_type) + + self.train_dataset = DLDataLoader.create( + framework=self.framework_type, + DLDataReader=self.dl1dh_reader, + indices=training_indices, + tasks=self.reco_tasks, + batch_size=self.batch_size, + random_seed=self.random_seed, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + parameters=self.parameters, + use_augmentation=self.use_augmentation, + is_training=True, + ) + self.training_loader = DataLoader( + dataset=self.train_dataset, + batch_size=None, + batch_sampler=None, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + prefetch_factor=4 if self.num_workers > 0 else None, + persistent_workers=self.persistent_workers if self.num_workers > 0 else False + ) + + print(len(self.training_loader)) + + self.validation_dataset = DLDataLoader.create( + framework=self.framework_type, + DLDataReader=self.dl1dh_reader, + indices=validation_indices, + tasks=self.reco_tasks, + batch_size=self.batch_size, + random_seed=self.random_seed, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + parameters=self.parameters, + use_augmentation=False, + is_training=False, + ) + self.validation_loader = DataLoader( + dataset=self.validation_dataset, + batch_size=None, + batch_sampler=None, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + prefetch_factor=4 if self.num_workers > 0 else None, + persistent_workers=self.persistent_workers if self.num_workers > 0 else False + ) + + print(len(self.validation_dataset)) + + + def start(self): + super().start() + for task in self.tasks: + + # Create the experiment folder + save_folder = create_experiment_folder( + f"run_{task.name}_training_", next_number=self.experiment_number + ) + + # Save the resolved configuration parameters to a config file in the experiment directory + config_filename = os.path.join(save_folder, "resolved_config.yml") + with open(config_filename, "w") as outfile: + import yaml + yaml.dump(self.parameters, outfile, default_flow_style=False) + + # ------------------------------------------------------------------------------ + # Select the model and precision + # ------------------------------------------------------------------------------ + + from ctlearn.core.pytorch.model import CTLearnPyTorchModel + + def load_pytorch_model_net(model_info, task_name, num_inputs, num_outputs): + model_name = model_info.get("model_name", "") + try: + component_cls = CTLearnPyTorchModel.non_abstract_subclasses().get(model_name) + if component_cls is not None: + params = model_info.get("parameters", {}).copy() + params.pop("task", None) + params.pop("num_inputs", None) + params.pop("num_outputs", None) + # Ensure device_str is set + params["parent"] = self + component = component_cls( + task=task_name, + num_inputs=num_inputs, + num_outputs=num_outputs, + **params + ) + return component.model + except Exception as e: + self.log.warning(f"Failed to load model {model_name} as Component: {e}. Falling back to create_model.") + return create_model(model_info) + + import torch.nn as nn + import torch + class ONNXModelWrapper(nn.Module): + def __init__(self, onnx_model_net, active_task, onnx_input_shape): + super().__init__() + self.onnx_model_net = onnx_model_net + self.active_task = active_task + self.onnx_input_shape = onnx_input_shape + + self.onnx_channel_last = False + self.expected_channels = 1 + if len(onnx_input_shape) == 4: + if onnx_input_shape[3] in [1, 2, 3] and onnx_input_shape[1] > onnx_input_shape[3]: + self.onnx_channel_last = True + self.expected_channels = onnx_input_shape[3] + else: + self.expected_channels = onnx_input_shape[1] + + def forward(self, x, y=None): + if self.expected_channels == 2 and y is not None: + x = torch.cat([x, y], dim=1) + + import inspect + sig = inspect.signature(self.onnx_model_net.forward) + num_onnx_inputs = len(sig.parameters) + + if num_onnx_inputs == 2 and y is not None: + if self.onnx_channel_last: + x = x.permute(0, 2, 3, 1) + y = y.permute(0, 2, 3, 1) + out = self.onnx_model_net(x, y) + else: + if self.onnx_channel_last: + x = x.permute(0, 2, 3, 1) + out = self.onnx_model_net(x) + + if isinstance(out, (tuple, list)): + if len(out) == 3: + return out + val = out[0] + else: + val = out + + if self.active_task == Task.type: + return val, None, None + elif self.active_task == Task.energy: + return None, val, None + else: + return None, None, val + + num_inputs = 1 ## Change thiss!!!! + + if self.load_onnx_model: + self.log.info(f"Loading ONNX model from {self.load_onnx_model} for training...") + self.log.warning( + "WARNING: Training an ONNX model from scratch (untrained) in PyTorch using onnx2pytorch " + "often leads to frozen gradients and the loss not improving. This is because onnx2pytorch " + "is designed primarily for inference, and many operations break the computational graph. " + "It is highly recommended to train the native PyTorch model instead by removing the " + "--load_onnx_model flag." + ) + import onnx + from onnx2pytorch import ConvertModel + try: + onnx_proto = onnx.load(self.load_onnx_model) + onnx_model = ConvertModel(onnx_proto) + onnx_input_shape = [dim.dim_value for dim in onnx_proto.graph.input[0].type.tensor_type.shape.dim] + model_net = ONNXModelWrapper(onnx_model, task, onnx_input_shape) + precision = self.parameters["arch"].get(f"precision_{task.name.lower()}", "32-true") + except Exception as e: + self.log.error(f"Failed to load ONNX model: {e}") + raise e + else: + if task == Task.type: + precision = self.parameters["arch"]["precision_type"] + model_net = load_pytorch_model_net(self.parameters["model"]["model_type"], "type", num_inputs, 2) + + elif task == Task.energy: + precision = self.parameters["arch"]["precision_energy"] + model_net = load_pytorch_model_net(self.parameters["model"]["model_energy"], "energy", num_inputs, 1) + + elif task == Task.cameradirection or task == Task.skydirection: + precision = self.parameters["arch"]["precision_direction"] + model_net = load_pytorch_model_net(self.parameters["model"]["model_direction"], "direction", num_inputs, 3) + + else: + raise ValueError( + f"task:{task.name} is not supported. Task must be type, direction or energy" + ) + + # if hasattr(model_net, 'T'): + # self.training_loader.set_T(model_net.T) + # ------------------------------------------------------------------------------ + # Load Checkpoints + # ------------------------------------------------------------------------------ + if task == Task.type: + check_point_path = self.parameters["data"]["type_checkpoint"] + + elif task == Task.energy: + check_point_path = self.parameters["data"]["energy_checkpoint"] + + elif task == Task.cameradirection or task == Task.skydirection: + check_point_path = self.parameters["data"]["direction_checkpoint"] + + else: + raise ValueError( + f"task:{task.name} is not supported. Task must be type, direction or energy" + ) + # Load the checkpoint if provided + if check_point_path: + model_net = ModelHelper.loadModel( + model_net, "", check_point_path, Mode.train, device_str=self.device_str + ) + + # Setup the TensorBoard logger + log_dir = save_folder + + tb_logger = TensorBoardLogger( + save_dir=log_dir, + name="exp_" + + str(self.experiment_number) + + "_" + + task.name + + "_train", + default_hp_metric=False, + ) + + extra_trainer_args = {} + test_limit = os.environ.get("CTLEARN_TEST_LIMIT") + if test_limit: + try: + limit = int(test_limit) + except ValueError: + limit = 5 + extra_trainer_args["limit_train_batches"] = limit + extra_trainer_args["limit_val_batches"] = limit + + # Setup the Trainer + trainer_pl = CTLearnTrainer( + max_epochs=self.parameters["hyp"]["epochs"], + accelerator=self.parameters["arch"]["device"], + devices=self.devices, + strategy= self.parameters["arch"]["strategy"], + default_root_dir=log_dir, + log_every_n_steps=1, + logger=tb_logger, + num_sanity_val_steps=0, + precision=precision, + gradient_clip_val=self.parameters["hyp"]["gradient_clip_val"], + callbacks=[GPUStatsLogger()], + sync_batchnorm=True, + enable_progress_bar=not self.disable_progress_bar, + **extra_trainer_args + ) + + # Setup Lighting + lightning_model = CTLearnPL( + model=model_net, + save_folder=trainer_pl.get_log_dir(), + task=task, + mode = Mode.train, + parameters=self.parameters, + k=self.save_k, + train_loader= self.training_loader, + val_loader= self.validation_loader, + train_dataset=self.train_dataset, + val_dataset=self.validation_dataset + ) + + if trainer_pl.is_global_zero: + # Save configuration file. + if not os.path.exists(trainer_pl.get_log_dir()): + os.makedirs(trainer_pl.get_log_dir()) + + with open(os.path.join(trainer_pl.get_log_dir(),"parameters.json"), "w") as f: + def path_serializer(obj): + import pathlib + if isinstance(obj, pathlib.Path): + return str(obj) + raise TypeError(f"Type {type(obj)} not serializable") + json.dump(self.parameters, f, indent=4, default=path_serializer) + + print(f"Run tensorboard server: tensorboard --load_fast=false --host=0.0.0.0 --logdir={trainer_pl.get_log_dir()}/") + + print(f"Accelerator: {trainer_pl.accelerator}") + print(f"Num. Devices: {trainer_pl.num_devices}") + + trainer_pl.fit( + model=lightning_model, + train_dataloaders=self.training_loader, + val_dataloaders=[self.validation_loader], + ) + + # Export to ONNX if requested + if self.save_onnx: + self.log.info("Converting PyTorch model into ONNX format...") + try: + # Load the best model weights if available + if trainer_pl.checkpoint_callback and trainer_pl.checkpoint_callback.best_model_path: + best_model_path = trainer_pl.checkpoint_callback.best_model_path + self.log.info(f"Loading best model from {best_model_path} for ONNX export...") + model_net = ModelHelper.loadModel( + model_net, "", best_model_path, Mode.train, device_str=self.device_str + ) + + # Create dummy input dynamically from a batch + batch = next(iter(self.training_loader)) + features, labels, t = batch + + import inspect + sig = inspect.signature(model_net.forward) + num_inputs_sig = len(sig.parameters) + + # Transfer dummy inputs to same device as model + if num_inputs_sig == 2: + dummy_image = torch.randn_like(features["image"][:1]).to(self.device_str) + dummy_peak = torch.randn_like(features["peak_time"][:1]).to(self.device_str) + dummy_input = (dummy_image, dummy_peak) + input_names = ["image", "peak_time"] + else: + dummy_input = torch.randn_like(features["image"][:1]).to(self.device_str) + input_names = ["image"] + + output_names = [task.name] + + # Put model in evaluation mode + model_net.eval() + + onnx_path = os.path.join(save_folder, f"ctlearn_model_{task.name}") + ModelHelper.exportOnnx( + model_net, + dummy_input, + onnx_path, + input_names=input_names, + output_names=output_names + ) + self.log.info(f"ONNX model saved successfully to {onnx_path}.onnx and {onnx_path}_simp.onnx") + except Exception as e: + self.log.error(f"Failed to export model to ONNX: {e}") + + + + def show_version(self): + print("Pytorch 2.3") diff --git a/ctlearn/tools/pytorch/utils.py b/ctlearn/tools/pytorch/utils.py new file mode 100644 index 00000000..4d10fe47 --- /dev/null +++ b/ctlearn/tools/pytorch/utils.py @@ -0,0 +1,190 @@ +from ctlearn.core.ctlearn_enum import Task +from typing import List +import os +import yaml + +expected_structure = { + "data": { + "train_gamma_proton": None, + "validation_gamma_proton": None, + "train_gamma": None, + "validation_gamma": None, + "test_gamma": None, + "test_proton": None, + "test_electron": None, + "test_validation_gamma": None, + "test_validation_gamma_proton": None, + }, + "run_details": { + "mode": None, + "task": None, + "test_type": None, + "experiment_number": None, + }, + "cut-off": {"leakage_intensity": None, "intensity": None}, + "model": { + "model_type": { + "model_name": None, + "parameters": None, + }, + "model_energy": { + "model_name": None, + "parameters": None, + }, + "model_direction": { + "model_name": None, + "parameters": None, + }, + }, + "hyp": { + "epochs": None, + "batches": None, + "dynamic_batches": None, + "optimizer": None, + "momentum": None, + "weight_decay": None, + "learning_rate": None, + "lrf": None, + "start_epoch": None, + "steps_epoch": None, + "l2_lambda": None, + "adam_epsilon": None, + "gradient_clip_val": None, + "save_k": None, + }, + "augmentation": { + "use_augmentation": None, + "aug_prob": None, + "rot_prob": None, + "trans_prob": None, + "flip_hor_prob": None, + "flip_ver_prob": None, + "mask_prob": None, + "noise_prob": None, + "max_rot": None, + "max_trans": None, + }, + "normalization": { + "use_clean": None, + "type_mu": None, + "type_sigma": None, + "dir_mu": None, + "dir_sigma": None, + "energy_mu": None, + "energy_sigma": None, + }, + "dataset": {"num_workers": None, "pin_memory": None, "persistent_workers": None}, + "arch": { + "device": None, + "precision_type": None, + "precision_energy": None, + "precision_direction": None, + "devices": None, + "strategy": None, + }, +} + + +# ------------------------------------------------------------------------------------------------------------------- +# Sanity check function +def sanity_check(config, expected_structure): + """ + Recursively checks if all required keys in the expected_structure exist in the config data. + Raises a KeyError if a key is missing. + """ + for key, substructure in expected_structure.items(): + if key not in config: + raise KeyError(f"Missing key: {key}") + + # If the substructure is a dictionary, recursively check the subkeys + if isinstance(substructure, dict): + if not isinstance(config[key], dict): + raise KeyError( + f"Expected a dictionary for key: {key}, but got: {type(config[key])}" + ) + sanity_check(config[key], substructure) + + +# ------------------------------------------------------------------------------------------------------------------- +def read_configuration(config_file_str="./config/training_config.yml"): + parameters = None + + if os.path.exists(config_file_str): + with open(config_file_str, "r") as config_file: + parameters = yaml.safe_load(config_file) + + else: + print(f"Configuration file not found. ({config_file_str})") + + return parameters + + +# ------------------------------------------------------------------------------------------------------------------- +def create_experiment_folder(prefix="run_", next_number=None): + """ + Create the next folder within the specified directory with a given prefix. + If next_number is not specified, automatically find the next available number. + + :param run_directory: The directory where folders are managed. + :param prefix: Prefix used for folders. + :param next_number: Optional. Specify the number to be used for the new folder. + """ + run_directory = "./run" + + # Ensure the 'run' directory exists + if not os.path.exists(run_directory): + os.makedirs(run_directory) + print(f"Directory '{run_directory}' created.") + + if next_number is None: + # List all subdirectories in the 'run' directory + folders = [ + f + for f in os.listdir(run_directory) + if os.path.isdir(os.path.join(run_directory, f)) + ] + # Filter folders that match the prefix pattern and end with a digit + matching_folders = [ + f for f in folders if f.startswith(prefix) and f[len(prefix) :].isdigit() + ] + + # Find the highest number and calculate the next one + if matching_folders: + highest_number = max( + int(folder[len(prefix) :]) for folder in matching_folders + ) + next_number = highest_number + 1 + else: + next_number = 0 + + # Create the new folder with the next number + new_folder_name = f"{prefix}{next_number}" + new_folder_path = os.path.join(run_directory, new_folder_name) + new_folder_path += "/" + if not os.path.exists(new_folder_path): + os.makedirs(new_folder_path) + print(f"New folder created: {new_folder_path}") + return new_folder_path + + +# ------------------------------------------------------------------------------------------------------------------- +def str_list_to_enum_list(reco_tasks: List) -> List[Task]: + + tasks = [] + for task_str in reco_tasks: + try: + tasks.append(Task[task_str]) + except KeyError: + print(f"'{task_str}' is not a valid enum type.") + return tasks + + +# ------------------------------------------------------------------------------------------------------------------- +def get_absolute_config_path(config_file_str="./ctlearn/tools/train/pytorch/config/default_config_file.yml") -> str: + """ + Returns the absolute path of the configuration file. + + :param config_file_str: Relative or absolute path to the configuration file. + :return: Absolute path of the configuration file. + """ + return os.path.abspath(config_file_str) \ No newline at end of file diff --git a/ctlearn/tools/tests/test_train_model.py b/ctlearn/tools/tests/test_train_model.py index 71f3f6fb..be4be06c 100644 --- a/ctlearn/tools/tests/test_train_model.py +++ b/ctlearn/tools/tests/test_train_model.py @@ -2,20 +2,22 @@ import pytest import shutil -pytest.importorskip("tensorflow") -from unittest import mock - from ctapipe.core import run_tool -from ctlearn.tools import DLFrameWork - +from ctlearn.conftest import TRAINING_TOOLS, MODEL_FILE_FORMATS +@pytest.mark.parametrize("framework", ["Keras"]) +@pytest.mark.parametrize("model", ["SingleCNN", "ResNet", "LoadedModel"]) @pytest.mark.parametrize("reco_task", ["type", "energy", "cameradirection"]) -@mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True) -def test_train_ctlearn_model(mock_eq, reco_task, dl1_gamma_file, dl1_proton_file, tmp_path): +def test_train_ctlearn_model(framework, model, reco_task, dl1_gamma_file, dl1_proton_file, ctlearn_trained_dl1_mono_models, tmp_path): """ Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks. Each test run gets its own isolated temp directories. """ + + # Restrict to MST array + telescope_type = "MST" + allowed_tels = [7, 13, 15, 16, 17, 19] + # Temporary directories for signal and background signal_dir = tmp_path / "gamma_dl1" signal_dir.mkdir(parents=True, exist_ok=True) @@ -28,9 +30,20 @@ def test_train_ctlearn_model(mock_eq, reco_task, dl1_gamma_file, dl1_proton_file # Hardcopy DL1 proton file to the background directory shutil.copy(dl1_proton_file, background_dir) + # Hardcopy the trained models to the model directory + if model == "LoadedModel": + model_dir = tmp_path / "pretrained_model" + model_dir.mkdir(parents=True, exist_ok=True) + key = f"{framework}_{telescope_type}_{reco_task}" + shutil.copy( + ctlearn_trained_dl1_mono_models[key], + model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}", + ) + model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" + assert model_file.exists(), f"Trained {framework} mono model file not found for {key}" + # Output directory for trained model - output_dir = tmp_path / f"ctlearn_{reco_task}" - allowed_tels = [7, 13, 15, 16, 17, 19] + output_dir = tmp_path / f"ctlearn_{framework}_{model}_{reco_task}" # Build command-line arguments argv = [ @@ -44,37 +57,41 @@ def test_train_ctlearn_model(mock_eq, reco_task, dl1_gamma_file, dl1_proton_file f"--DLImageReader.allowed_tels={allowed_tels}", ] + # Include background only for classification task if reco_task == "type": argv.extend( [ f"--background={background_dir}", "--pattern-background=*.dl1.h5", + "--DLImageReader.enforce_subarray_equality=False", ] ) - # Run training - assert run_tool(DLFrameWork(), argv=argv, cwd=tmp_path) == 0 - + argv.append(f"--TrainCTLearnModel.model_type={model}") + if model == "LoadedModel": + argv.append(f"--LoadedModel.load_model_from={model_file}") + assert run_tool(TRAINING_TOOLS[framework](), argv=argv, cwd=tmp_path) == 0 # --- Additional checks --- # Check that the trained model exists - model_file = output_dir / "ctlearn_model.keras" + model_file = output_dir / f"ctlearn_model.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained model file not found for {reco_task}" - # Check training_log.csv exists - log_file = output_dir / "training_log.csv" - assert log_file.exists(), f"Training log file not found for {reco_task}" - # Read CSV and verify number of epochs - log_df = pd.read_csv(log_file) - num_epochs_logged = log_df.shape[0] - assert ( - num_epochs_logged == 2 - ), f"Expected two epochs, found {num_epochs_logged} for {reco_task}" - # Check that val_loss column exists - assert ( - "val_loss" in log_df.columns - ), f"'val_loss' column missing in training_log.csv for {reco_task}" - val_loss = log_df["val_loss"].dropna() - assert not val_loss.empty, f"'val_loss' column is empty for {reco_task}" - assert ((val_loss >= 0.0) & (val_loss <= 1.0)).all(), ( - f"'val_loss' values out of range [0.0, 1.0] for {reco_task}: " - f"{val_loss.tolist()}" - ) + if framework == "Keras": + # Check training_log.csv exists + log_file = output_dir / "training_log.csv" + assert log_file.exists(), f"Training log file not found for {reco_task}" + # Read CSV and verify number of epochs + log_df = pd.read_csv(log_file) + num_epochs_logged = log_df.shape[0] + assert ( + num_epochs_logged == 2 + ), f"Expected two epochs, found {num_epochs_logged} for {reco_task}" + # Check that val_loss column exists + assert ( + "val_loss" in log_df.columns + ), f"'val_loss' column missing in training_log.csv for {reco_task}" + val_loss = log_df["val_loss"].dropna() + assert not val_loss.empty, f"'val_loss' column is empty for {reco_task}" + assert ((val_loss >= 0.0) & (val_loss <= 1.0)).all(), ( + f"'val_loss' values out of range [0.0, 1.0] for {reco_task}: " + f"{val_loss.tolist()}" + ) \ No newline at end of file diff --git a/ctlearn/tools/train_model.py b/ctlearn/tools/train_model.py index 19be869d..ec16cbfa 100644 --- a/ctlearn/tools/train_model.py +++ b/ctlearn/tools/train_model.py @@ -1,274 +1,260 @@ -import sys -import os -import warnings +""" +Base tool to train a ``CTLearnModel``on R1/DL1a data using the ``DLDataReader`` and ``DLDataLoader``. +""" -is_debug = '--debug' in sys.argv or any(arg.startswith('--log-level=DEBUG') for arg in sys.argv) -if not is_debug: - os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' - os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' - os.environ['NCCL_DEBUG'] = 'WARN' - warnings.filterwarnings("ignore", category=UserWarning) - warnings.filterwarnings("ignore", ".*NoneDefaultNotAllowedWarning.*") - warnings.filterwarnings("ignore", ".*MergeConflictWarning.*") - warnings.filterwarnings("ignore", ".*'ctlearn.tools.train_model' found in sys.modules.*") - -import atexit -import pandas as pd +from abc import abstractmethod import numpy as np + from ctapipe.core import Tool -from ctapipe.core.traits import CaselessStrEnum, Dict -from ctlearn.core.ctlearn_enum import FrameworkType -class DLFrameWork(Tool): +from ctapipe.core.tool import ToolConfigurationError +from ctapipe.core.traits import ( + Bool, + CaselessStrEnum, + Path, + Float, + Int, + List, + Dict, + classes_with_traits, + ComponentName, + Unicode, +) +from ctlearn import __version__ as ctlearn_version +from ctlearn.core.model import CTLearnModel +from dl1_data_handler.reader import DLDataReader + + +class TrainCTLearnModel(Tool): """ - Tool to select and run a specific deep learning training framework (Keras or PyTorch) - for CTLearn model training. It dynamically loads the appropriate subclass based on - the user-defined --framework argument. + Base tool to train a ``~ctlearn.core.model.CTLearnModel`` on R1/DL1a data. + + The tool holds configurations and set up functions. + """ - name = "dlframework" + input_dir_signal = Path( + help="Input directory for signal events", + allow_none=False, + exists=True, + directory_ok=True, + file_ok=False, + ).tag(config=True) - framework_type = CaselessStrEnum( - ["pytorch", "keras"], - default_value="keras", - help="Framework to use: pytorch or keras", + file_pattern_signal = List( + trait=Unicode(), + default_value=["*.h5"], + help="List of specific file pattern for matching files in ``input_dir_signal``", ).tag(config=True) - early_stopping = Dict( + input_dir_background = Path( default_value=None, + help="Input directory for background events", allow_none=True, + exists=True, + directory_ok=True, + file_ok=False, + ).tag(config=True) + + file_pattern_background = List( + trait=Unicode(), + default_value=["*.h5"], + help="List of specific file pattern for matching files in ``input_dir_background``", + ).tag(config=True) + + dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( + config=True + ) + + stack_telescope_images = Bool( + default_value=False, + allow_none=False, help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + "Set whether to stack the telescope images in the data loader. " + "Requires DLDataReader mode to be ``stereo``." ), ).tag(config=True) - early_stopping = Dict( - default_value=None, + sort_by_intensity = Bool( + default_value=False, allow_none=True, help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + "Set whether to sort the telescope images by intensity in the data loader. " + "Requires DLDataReader mode to be ``stereo``." ), ).tag(config=True) - early_stopping = Dict( - default_value=None, - allow_none=True, + model_type = CaselessStrEnum( + ["SingleCNN", "ResNet", "LoadedModel"], + default_value="ResNet", + allow_none=False, help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + "Model type to be used in the Keras or PyTorch framework. " + "The framework is determined by the inherited tools being used." ), ).tag(config=True) - early_stopping = Dict( + output_dir = Path( + exits=False, default_value=None, - allow_none=True, + allow_none=False, + directory_ok=True, + file_ok=False, + help="Output directory for the trained reconstructor.", + ).tag(config=True) + + reco_tasks = List( + trait=CaselessStrEnum(["type", "energy", "cameradirection", "skydirection"]), + allow_none=False, help=( - "Early stopping parameters for the Keras callback. " - "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + "List of reconstruction tasks to perform. " + "'type': classification of the primary particle type; " + "'energy': regression of the primary particle energy; " + "'cameradirection': reconstruction of the primary particle arrival direction in camera coordinates; " + "'skydirection': reconstruction of the primary particle arrival direction in sky coordinates." ), ).tag(config=True) - aliases = { - "framework": "DLFrameWork.framework_type", - } - - try: - from ctlearn.tools.train.pytorch.train_pytorch_model import TrainPyTorchModel - aliases.update(TrainPyTorchModel.aliases) - except ImportError: - pass - - try: - from ctlearn.tools.train.keras.train_keras_model import TrainKerasModel - aliases.update(TrainKerasModel.aliases) - except ImportError: - pass - - def __init__(self, **kwargs): - """ - Initialize the DLFrameWork tool and prepare for framework injection. - """ - super().__init__(**kwargs) - self.framework_instance = None - - def setup(self): - """ - Setup method called after basic trait parsing. - This dynamically loads and prepares the correct framework subclass - (TrainKerasModel or TrainPyTorchModel). - """ - framework_enum = self.string_to_type(self.framework_type) - self.framework_instance = self.get_framework(framework_enum) - - # Inject aliases and shared config before full CLI parsing - self.framework_instance.update_config(self.config) - self.aliases.update(self.framework_instance.aliases) - DLFrameWork.aliases.update(self.framework_instance.aliases) - - def start(self): - """ - Start method called after setup. Executes the selected framework instance. - """ - print("start") - self.framework_instance.run() - - @classmethod - def string_to_type(cls, str_type: str) -> FrameworkType: - """ - Convert a string to a FrameworkType enum (case-insensitive). - - Parameters: - str_type (str): The name of the framework (e.g., 'keras', 'pytorch'). - - Returns: - FrameworkType: Corresponding enum value. + n_epochs = Int( + default_value=10, + allow_none=False, + help="Number of epochs to train the neural network.", + ).tag(config=True) - Raises: - ValueError: If the provided string is not a valid framework type. - """ - try: - return FrameworkType[str_type.upper()] - except KeyError: - raise ValueError(f"'{str_type}' is not a valid framework type.") + batch_size = Int( + default_value=64, + allow_none=False, + help="Size of the batch to train the neural network.", + ).tag(config=True) - @classmethod - def get_framework(cls, framework_type: FrameworkType): - """ - Dynamically import and return the corresponding training class - based on the framework type. + validation_split = Float( + default_value=0.1, + help="Fraction of the data to use for validation", + min=0.01, + max=0.99, + ).tag(config=True) - Parameters: - framework_type (FrameworkType): Enum indicating which framework to use. + optimizer = Dict( + default_value={ + "name": "Adam", + "base_learning_rate": 0.0001, + "adam_epsilon": 1.0e-8, + }, + help=( + "Optimizer to use for training. " + "E.g. {'name': 'Adam', 'base_learning_rate': 0.0001, 'adam_epsilon': 1.0e-8}. " + ), + ).tag(config=True) - Returns: - Tool: An instance of the selected training framework (subclass of Tool). + random_seed = Int( + default_value=0, + help=( + "Random seed for shuffling the data " + "before the training/validation split " + "and after the end of an epoch." + ), + ).tag(config=True) - Raises: - ImportError: If the training module could not be imported. - ValueError: If the framework type is unknown. - """ - if framework_type == FrameworkType.KERAS: - try: - from ctlearn.tools.train.keras.train_keras_model import TrainKerasModel + aliases = { + "signal": "TrainCTLearnModel.input_dir_signal", + "background": "TrainCTLearnModel.input_dir_background", + "pattern-signal": "TrainCTLearnModel.file_pattern_signal", + "pattern-background": "TrainCTLearnModel.file_pattern_background", + "reco": "TrainCTLearnModel.reco_tasks", + ("o", "output"): "TrainCTLearnModel.output_dir", + } - fw = TrainKerasModel() - except ImportError as e: - raise ImportError(f"Not possible to import TrainKerasModel: {e}") from e + classes = classes_with_traits(CTLearnModel) + classes_with_traits(DLDataReader) - elif framework_type == FrameworkType.PYTORCH: - try: - from ctlearn.tools.train.pytorch.train_pytorch_model import ( - TrainPyTorchModel, + def setup(self): + self.log.info("ctlearn version %s", ctlearn_version) + # Check if the output directory exists + if self.output_dir.exists(): + raise ToolConfigurationError( + f"Output directory {self.output_dir} already exists." + ) + # Get signal input files + self.input_url_signal = [] + for signal_pattern in self.file_pattern_signal: + self.input_url_signal.extend(self.input_dir_signal.glob(signal_pattern)) + # Get bkg input files + self.input_url_background = [] + if self.input_dir_background is not None: + for background_pattern in self.file_pattern_background: + self.input_url_background.extend( + self.input_dir_background.glob(background_pattern) ) - fw = TrainPyTorchModel() - - except ImportError as e: - raise ImportError( - f"Not possible to import TrainPyTorchModel: {e}" - ) from e - - else: - raise ValueError(f"Unknown Framework: {framework_type.name}") - - return fw - - from ctlearn.tools.train.base_train_model import TrainCTLearnModel - classes = [TrainCTLearnModel] - try: - from ctlearn.tools.train.keras.train_keras_model import TrainKerasModel - classes.append(TrainKerasModel) - except ImportError: - pass - - try: - from ctlearn.tools.train.pytorch.train_pytorch_model import TrainPyTorchModel - classes.append(TrainPyTorchModel) - except ImportError: + # Set up the data reader + self.log.info("Loading data:") + self.log.info("For a large dataset, this may take a while...") + if self.dl1dh_reader_type == "DLFeatureVectorReader": + raise NotImplementedError( + "'DLFeatureVectorReader' is not supported in CTLearn yet. " + "Missing stereo CTLearnModel implementation." + ) + self.dl1dh_reader = DLDataReader.from_name( + self.dl1dh_reader_type, + input_url_signal=sorted(self.input_url_signal), + input_url_background=sorted(self.input_url_background), + parent=self, + ) + self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) + if "type" in self.reco_tasks: + self.log.info( + "Number of signal events: %d", self.dl1dh_reader.n_signal_events + ) + self.log.info( + "Number of background events: %d", self.dl1dh_reader.n_bkg_events + ) + # Check if the number of events is enough to form a batch + if self.dl1dh_reader._get_n_events() < self.batch_size: + raise ValueError( + f"{self.dl1dh_reader._get_n_events()} events are not enough " + f"to form a batch of size {self.batch_size}. Reduce the batch size." + ) + # Check if there are at least two classes in the reader for the particle classification + if self.dl1dh_reader.class_weight is None and "type" in self.reco_tasks: + raise ValueError( + "Classification task selected but less than two classes are present in the data." + ) + # Check if stereo mode is selected for stacking telescope images + if self.stack_telescope_images and self.dl1dh_reader.mode == "mono": + raise ToolConfigurationError( + f"Cannot stack telescope images in mono mode. Use stereo mode for stacking." + ) + # Ckeck if only one telescope type is selected for stacking telescope images + if ( + self.stack_telescope_images + and len(list(self.dl1dh_reader.selected_telescopes)) > 1 + ): + raise ToolConfigurationError( + f"Cannot stack telescope images from multiple telescope types. Use only one telescope type." + ) + # Check if sorting by intensity is disabled for stacking telescope images + if self.stack_telescope_images and self.sort_by_intensity: + raise ToolConfigurationError( + f"Cannot stack telescope images when sorting by intensity. Disable sorting by intensity." + ) + + # Set up the data loaders for training and validation + self.indices = list(range(self.dl1dh_reader._get_n_events())) + # Shuffle the indices before the training/validation split + np.random.seed(self.random_seed) + np.random.shuffle(self.indices) + self.n_validation_examples = int( + self.validation_split * self.dl1dh_reader._get_n_events() + ) + self.training_indices = self.indices[self.n_validation_examples:] + self.validation_indices = self.indices[:self.n_validation_examples] + + # Set up framework-specific training tool + self.setup_framework() + + @abstractmethod + def setup_framework(): + """ This is an abstract method for the setup of the framework-specific training tool.""" pass - - from ctapipe.core.traits import classes_with_traits - from dl1_data_handler.reader import DLDataReader - classes = classes + classes_with_traits(DLDataReader) - -def main(): - # Run the tool - tool = DLFrameWork() - - # Manually parse --framework to determine which subclass to load, as traitlets alias - # update can sometimes fail to parse it correctly before setup - framework = "keras" - for i, arg in enumerate(sys.argv[1:]): - if arg.startswith("--framework="): - framework = arg.split("=")[1].strip() - elif arg == "--framework" and i + 2 < len(sys.argv): - framework = sys.argv[i + 2].strip() - - tool.framework_type = framework - - minimal_args = [ - arg for arg in sys.argv[1:] if "--framework" in arg or arg in ["-h", "--help"] - ] - tool.initialize(argv=minimal_args) - - # Setup and inject the correct framework instance - tool.setup() - - # Parse all CLI args with the selected framework subclass - tool.framework_instance.initialize(argv=sys.argv[1:]) - - tool.run() - - -if __name__ == "__main__": - main() - - -# Example: -# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir2 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite -# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_*.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --pattern-background proton_theta_23.161_az_99.261_runs833-1250*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 - -# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 - - - - -# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 - - -# --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 - - - -# Type -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & - -# Direction -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & - - - -# Energy -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_1.yml> nohup_energy_training.out 2>&1 & - - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs243-302.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs303-362.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs363-422.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs423-482.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs483-541.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs542-600.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & - -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_6.000_az_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & -# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_2.yml> nohup_energy_training.out 2>&1 & \ No newline at end of file + def finish(self): + self.log.info("Tool is shutting down") diff --git a/ctlearn/tools/train_model_current.py b/ctlearn/tools/train_model_current.py new file mode 100644 index 00000000..84b2b264 --- /dev/null +++ b/ctlearn/tools/train_model_current.py @@ -0,0 +1,267 @@ +import atexit +import pandas as pd +import numpy as np +import sys +from ctapipe.core import Tool +from ctapipe.core.traits import CaselessStrEnum, Dict +from ctlearn.core.ctlearn_enum import FrameworkType +class DLFrameWork(Tool): + """ + Tool to select and run a specific deep learning training framework (Keras or PyTorch) + for CTLearn model training. It dynamically loads the appropriate subclass based on + the user-defined --framework argument. + """ + + name = "dlframework" + + framework_type = CaselessStrEnum( + ["pytorch", "keras"], + default_value="keras", + help="Framework to use: pytorch or keras", + ).tag(config=True) + + early_stopping = Dict( + default_value=None, + allow_none=True, + help=( + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + ), + ).tag(config=True) + + early_stopping = Dict( + default_value=None, + allow_none=True, + help=( + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + ), + ).tag(config=True) + + early_stopping = Dict( + default_value=None, + allow_none=True, + help=( + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + ), + ).tag(config=True) + + early_stopping = Dict( + default_value=None, + allow_none=True, + help=( + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + ), + ).tag(config=True) + + aliases = { + "framework": "DLFrameWork.framework_type", + } + + try: + from pytorch_ctlearn.ctlearn.ctlearn.tools.pytorch.train_pytorch_model import TrainPyTorchModel + aliases.update(TrainPyTorchModel.aliases) + except ImportError: + pass + + try: + from pytorch_ctlearn.ctlearn.ctlearn.tools.keras.train_model import TrainKerasModel + aliases.update(TrainKerasModel.aliases) + except ImportError: + pass + + def __init__(self, **kwargs): + """ + Initialize the DLFrameWork tool and prepare for framework injection. + """ + super().__init__(**kwargs) + self.framework_instance = None + + def setup(self): + """ + Setup method called after basic trait parsing. + This dynamically loads and prepares the correct framework subclass + (TrainKerasModel or TrainPyTorchModel). + """ + framework_enum = self.string_to_type(self.framework_type) + self.framework_instance = self.get_framework(framework_enum) + + # Inject aliases and shared config before full CLI parsing + self.framework_instance.update_config(self.config) + self.aliases.update(self.framework_instance.aliases) + DLFrameWork.aliases.update(self.framework_instance.aliases) + + def start(self): + """ + Start method called after setup. Executes the selected framework instance. + """ + print("start") + self.framework_instance.run() + + @classmethod + def string_to_type(cls, str_type: str) -> FrameworkType: + """ + Convert a string to a FrameworkType enum (case-insensitive). + + Parameters: + str_type (str): The name of the framework (e.g., 'keras', 'pytorch'). + + Returns: + FrameworkType: Corresponding enum value. + + Raises: + ValueError: If the provided string is not a valid framework type. + """ + try: + return FrameworkType[str_type.upper()] + except KeyError: + raise ValueError(f"'{str_type}' is not a valid framework type.") + + @classmethod + def get_framework(cls, framework_type: FrameworkType): + """ + Dynamically import and return the corresponding training class + based on the framework type. + + Parameters: + framework_type (FrameworkType): Enum indicating which framework to use. + + Returns: + Tool: An instance of the selected training framework (subclass of Tool). + + Raises: + ImportError: If the training module could not be imported. + ValueError: If the framework type is unknown. + """ + if framework_type == FrameworkType.KERAS: + try: + from pytorch_ctlearn.ctlearn.ctlearn.tools.keras.train_model import TrainKerasModel + + fw = TrainKerasModel() + except ImportError as e: + raise ImportError(f"Not possible to import TrainKerasModel: {e}") from e + + elif framework_type == FrameworkType.PYTORCH: + try: + from pytorch_ctlearn.ctlearn.ctlearn.tools.pytorch.train_pytorch_model import ( + TrainPyTorchModel, + ) + + fw = TrainPyTorchModel() + + except ImportError as e: + raise ImportError( + f"Not possible to import TrainPyTorchModel: {e}" + ) from e + + else: + raise ValueError(f"Unknown Framework: {framework_type.name}") + + return fw + + @property + def classes(self): + from ctlearn.tools.train.base_train_model import TrainCTLearnModel + from ctapipe.core.traits import classes_with_traits + from dl1_data_handler.reader import DLDataReader + + tool_classes = [ + type(self), + TrainCTLearnModel, + ] + + try: + from pytorch_ctlearn.ctlearn.ctlearn.tools.keras.train_model import TrainKerasModel + tool_classes.append(TrainKerasModel) + except ImportError: + pass + + try: + from pytorch_ctlearn.ctlearn.ctlearn.tools.pytorch.train_pytorch_model import TrainPyTorchModel + tool_classes.append(TrainPyTorchModel) + except ImportError: + pass + + return tool_classes + classes_with_traits(DLDataReader) + +def main(): + # Run the tool + tool = DLFrameWork() + + # Manually parse --framework to determine which subclass to load, as traitlets alias + # update can sometimes fail to parse it correctly before setup + framework = "keras" + for i, arg in enumerate(sys.argv[1:]): + if arg.startswith("--framework="): + framework = arg.split("=")[1].strip() + elif arg == "--framework" and i + 2 < len(sys.argv): + framework = sys.argv[i + 2].strip() + + tool.framework_type = framework + + minimal_args = [ + arg for arg in sys.argv[1:] if "--framework" in arg or arg in ["-h", "--help"] + ] + tool.initialize(argv=minimal_args) + + # Setup and inject the correct framework instance + tool.setup() + + # Parse all CLI args with the selected framework subclass + tool.framework_instance.initialize(argv=sys.argv[1:]) + + tool.run() + + +if __name__ == "__main__": + main() + + +# Example: +# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir2 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite +# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_*.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --pattern-background proton_theta_23.161_az_99.261_runs833-1250*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 + +# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 + + + + +# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 + + +# --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 + + + +# Type +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# Direction +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & + + + +# Energy +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_1.yml> nohup_energy_training.out 2>&1 & + + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs243-302.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs303-362.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs363-422.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs423-482.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs483-541.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs542-600.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_6.000_az_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_2.yml> nohup_energy_training.out 2>&1 & \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index c6e155ba..27a567b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,7 @@ repository = "https://github.com/ctlearn-project/ctlearn" documentation = "https://ctlearn.readthedocs.io/en/latest/" [project.scripts] -ctlearn-train-model = "ctlearn.tools.train_model:main" +ctlearn-train-keras-model = "ctlearn.tools.keras.train_model:main" ctlearn-predict-mono-model = "ctlearn.tools.predict.predict_mono:main" ctlearn-predict-stereo-model = "ctlearn.tools.predict_stereo:main" ctlearn-predict-LST1= "ctlearn.tools.predict_LST1:main"