diff --git a/.flake8 b/.flake8 new file mode 100644 index 00000000..0baba3d4 --- /dev/null +++ b/.flake8 @@ -0,0 +1,3 @@ +[flake8] +max-line-length = 88 +extend-ignore = E203, W503 \ No newline at end of file diff --git a/.github/workflows/python-package-conda.yml b/.github/workflows/python-package-conda.yml index bea20d02..e4171bbe 100644 --- a/.github/workflows/python-package-conda.yml +++ b/.github/workflows/python-package-conda.yml @@ -2,12 +2,12 @@ name: CI on: push: - branches: - - "**" - tags: - - "**" + branches: ["**"] + tags: ["**"] pull_request: workflow_dispatch: + schedule: + - cron: "0 2 * * *" # Daily at 02:00 UTC jobs: build: @@ -17,14 +17,22 @@ jobs: python-version: ['3.12', '3.13', '3.14'] dl1dh-version: ['latest', 'nightly'] tensorflow-version: ['latest', '2.16.*'] + torch-version: ['latest', '2.4.*'] exclude: - python-version: '3.13' tensorflow-version: '2.16.*' - python-version: '3.14' tensorflow-version: '2.16.*' + - python-version: '3.13' + torch-version: '2.4.*' + - python-version: '3.14' + torch-version: '2.4.*' max-parallel: 6 runs-on: ${{ matrix.os }} continue-on-error: ${{ matrix.dl1dh-version == 'nightly' || matrix.python-version == '3.14' }} + env: + PIP_NO_CACHE_DIR: "1" + PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cpu" steps: - uses: actions/checkout@v4 @@ -61,6 +69,11 @@ jobs: else pip install "tensorflow==${{ matrix.tensorflow-version }}" fi + if [ "${{ matrix.torch-version }}" = "latest" ]; then + pip install --upgrade torch torchvision + else + pip install "torch==${{ matrix.torch-version }}" torchvision + fi - name: Add MKL_THREADING_LAYER variable run: echo "MKL_THREADING_LAYER=GNU" >> $GITHUB_ENV @@ -76,8 +89,8 @@ jobs: run: | source $HOME/miniconda/etc/profile.d/conda.sh conda activate ctlearn - pip install -e . - + pip install -e .[tests] + - name: Run pytest run: | source $HOME/miniconda/etc/profile.d/conda.sh diff --git a/.gitignore b/.gitignore index bd1fd54e..a45a88a4 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ ctlearn/_version.py - +output_dir2/ *.swp *.swo *.gemini @@ -12,13 +12,34 @@ ctlearn/_version.py *.png *.csv *~ +*.log +.vscode/ +launcher_workspace/ +*.h5 +output_dir2/ +output_dir/ +.vscode/ # Compiled Python files __pycache__/ *.py[cod] .DS_Store *.egg-info/ dist - +.pytest_cache/*.log # Sphinx documentation docs/build/ +build/ +# Default pytorch output +run/ +test/ +*.se2 +*.jar +*.puml +mc_tjark/ +calibration/ +test_local_cristian/ + +test/prepare_file.py +test_local_cristian/ +*.txt \ No newline at end of file diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 71fb5123..c4fb4307 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -7,9 +7,10 @@ build: os: ubuntu-22.04 tools: python: "3.12" - + jobs: + post_install: + # Installs ctlearn without downloading heavy PyTorch/TensorFlow binaries + - pip install --no-deps . python: install: - requirements: docs/requirements.txt - - method: pip - path: . \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index 4a08b151..cbbf0c70 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,11 +2,20 @@ FROM python:3.12 AS builder # Install git (needed for setuptools_scm during build) and build tool + RUN apt-get update \ && apt-get install -y --no-install-recommends git \ && rm -rf /var/lib/apt/lists/* RUN pip install --no-cache-dir build +# Copy source code needed for the build +WORKDIR /repo +COPY ./pyproject.toml MANIFEST.in ./ +COPY ./ctlearn ./ctlearn/ +# If .git is truly needed for versioning by setuptools_scm, copy it. Otherwise, omit. +COPY ./.git ./.git/ +RUN pip install --no-cache-dir build + # Copy source code needed for the build WORKDIR /repo COPY ./pyproject.toml MANIFEST.in ./ @@ -18,14 +27,22 @@ COPY ./.git ./.git/ RUN python -m build --wheel # Stage 2: Create the final runtime image BASED ON NVIDIA's TF image +# TODO what version to use ? after 24.?? TF 2.14 is not found in the container +FROM nvcr.io/nvidia/tensorflow:24.01-tf2-py3 + +# Copy only the built wheel from the builder stage's dist directory +# Build the wheel + +# Stage 2: Create the final runtime image BASED ON NVIDIA's TF image + FROM nvcr.io/nvidia/tensorflow:25.02-tf2-py3 # Copy only the built wheel from the builder stage's dist directory COPY --from=builder /repo/dist /tmp/dist +# Install the ctlearn wheel using pip from the NVIDIA base image # Install the ctlearn wheel using pip from the NVIDIA base image RUN python -m pip install --no-cache-dir /tmp/dist/* \ && rm -r /tmp/dist RUN addgroup --system ctlearn && adduser --system --group ctlearn -USER ctlearn - +USER ctlearn \ No newline at end of file diff --git a/README.rst b/README.rst index c5c4853f..92148a5c 100644 --- a/README.rst +++ b/README.rst @@ -47,6 +47,7 @@ The lastest version fo this package can be installed as a pip package: See the documentation for further information like `installation instructions for the IT-cluster `_, `installation instructions for developers `_, `package usage `_, and `dependencies `_ among other topics. + Citing this software -------------------- diff --git a/ctlearn/__init__.py b/ctlearn/__init__.py index c4c2735c..41ea362c 100644 --- a/ctlearn/__init__.py +++ b/ctlearn/__init__.py @@ -1,3 +1,3 @@ -from ._version import __version__ - -__all__ = ["__version__"] +from ._version import __version__ + +__all__ = ["__version__"] \ No newline at end of file diff --git a/ctlearn/conftest.py b/ctlearn/conftest.py index cdfcdb18..803189a3 100644 --- a/ctlearn/conftest.py +++ b/ctlearn/conftest.py @@ -7,16 +7,20 @@ import numpy as np import pytest import shutil -from astropy import units as u from astropy.table import Column, Table from traitlets.config.loader import Config from ctapipe.core import run_tool from ctapipe.io import write_table +from ctapipe.tools.process import ProcessorTool from ctapipe.utils import get_dataset_path -from ctlearn.tools import TrainCTLearnModel -from ctlearn.utils import get_lst1_subarray_description +from ctlearn.tools.keras.train_model import TrainCTLearnKerasModel +from ctlearn.tools.pytorch.train_model import TrainCTLearnPyTorchModel +from ctlearn.tools.utils import get_lst1_subarray_description + +TRAINING_TOOLS = {"Keras": TrainCTLearnKerasModel, "PyTorch": TrainCTLearnPyTorchModel} +MODEL_FILE_FORMATS = {"Keras": "keras", "PyTorch": "pth"} @pytest.fixture(scope="session") def gamma_simtel_path(): @@ -133,8 +137,6 @@ def dl1_gamma_file(dl1_tmp_path, gamma_simtel_path): """ DL1 file containing both images and parameters from a gamma simulation set. """ - from ctapipe.tools.process import ProcessorTool - output = dl1_tmp_path / "gamma.dl1.h5" argv = [ f"--input={gamma_simtel_path}", @@ -169,10 +171,7 @@ def r1_gamma_file(r1_tmp_path, gamma_simtel_path): """ R1 file containing both waveforms and parameters from a gamma simulation set. """ - from ctapipe.tools.process import ProcessorTool - output = r1_tmp_path / "gamma.r1.h5" - allowed_tels = [1, 2] argv = [ f"--input={gamma_simtel_path}", @@ -190,8 +189,6 @@ def r1_proton_file(r1_tmp_path, proton_simtel_path): """ R1 file containing both waveforms and parameters from a proton simulation set. """ - from ctapipe.tools.process import ProcessorTool - # Restrict to two LSTs for R1 tests to reduce computational load allowed_tels = [1, 2] output = r1_tmp_path / "proton.r1.h5" @@ -241,14 +238,10 @@ def ctlearn_trained_r1_mono_models(r1_gamma_file, r1_proton_file, tmp_path_facto # Loop over reconstruction tasks and train models for each combination ctlearn_trained_r1_mono_models = {} 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", @@ -267,14 +260,17 @@ def ctlearn_trained_r1_mono_models(r1_gamma_file, r1_proton_file, tmp_path_facto ] ) - # Run training - assert run_tool(TrainCTLearnModel(config=config), argv=argv, cwd=tmp_path) == 0 - - ctlearn_trained_r1_mono_models[f"{telescope_type}_{reco_task}"] = ( - output_dir / "ctlearn_model.keras" - ) - # Check that the trained model exists - assert ctlearn_trained_r1_mono_models[f"{telescope_type}_{reco_task}"].exists() + # Run training tools + for framework, training_tool in TRAINING_TOOLS.items(): + framework_argv = argv.copy() + output_dir = tmp_path / f"ctlearn_{framework}_{telescope_type}_{reco_task}" + framework_argv.append(f"--output={output_dir}") + assert run_tool(training_tool(config=config), argv=framework_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"{framework}_{telescope_type}_{reco_task}"].exists() return ctlearn_trained_r1_mono_models @@ -324,21 +320,16 @@ def ctlearn_trained_dl1_mono_models(dl1_gamma_file, dl1_proton_file, tmp_path_fa ctlearn_trained_dl1_mono_models = {} 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( @@ -349,20 +340,20 @@ def ctlearn_trained_dl1_mono_models(dl1_gamma_file, dl1_proton_file, tmp_path_fa f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", ] ) - - # Run training - assert ( - run_tool(TrainCTLearnModel(config=config), argv=argv, cwd=tmp_path) == 0 - ) - - ctlearn_trained_dl1_mono_models[f"{telescope_type}_{reco_task}"] = ( - output_dir / "ctlearn_model.keras" - ) - # 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(): + framework_argv = argv.copy() + output_dir = tmp_path / f"ctlearn_{framework}_{telescope_type}_{reco_task}" + framework_argv.append(f"--output={output_dir}") + assert run_tool(training_tool(config=config), argv=framework_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,14 +395,10 @@ def ctlearn_trained_dl1_stereo_models( # Loop over reconstruction tasks and train models for each combination ctlearn_trained_dl1_stereo_models = {} 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", @@ -431,14 +418,17 @@ def ctlearn_trained_dl1_stereo_models( ] ) - # Run training - assert run_tool(TrainCTLearnModel(config=config), argv=argv, cwd=tmp_path) == 0 - - ctlearn_trained_dl1_stereo_models[f"{telescope_type}_{reco_task}"] = ( - output_dir / "ctlearn_model.keras" - ) - # Check that the trained model exists - assert ctlearn_trained_dl1_stereo_models[ - f"{telescope_type}_{reco_task}" - ].exists() - return ctlearn_trained_dl1_stereo_models + # Run training tools + for framework, training_tool in TRAINING_TOOLS.items(): + framework_argv = argv.copy() + output_dir = tmp_path / f"ctlearn_{framework}_{telescope_type}_{reco_task}" + framework_argv.append(f"--output={output_dir}") + assert run_tool(training_tool(config=config), argv=framework_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"{framework}_{telescope_type}_{reco_task}" + ].exists() + return ctlearn_trained_dl1_stereo_models \ No newline at end of file diff --git a/ctlearn/core/__init__.py b/ctlearn/core/__init__.py new file mode 100644 index 00000000..6f68d8f3 --- /dev/null +++ b/ctlearn/core/__init__.py @@ -0,0 +1,64 @@ +""" +ctlearn core functionalities +""" + +from .model import ( + CTLearnModel, + SingleCNN, + ResNet, + LoadedModel, +) +from .keras.model import ( + build_fully_connect_keras_head, + KerasSingleCNN, + KerasResNet, + KerasLoadedModel, +) +from .keras.attention import ( + dual_squeeze_excite_block, + channel_squeeze_excite_block, + spatial_squeeze_excite_block, +) +from .keras.sequence import KerasSequence +from .pytorch.model import ( + BasicBlock, + BottleneckBlock, + MultiFullyConnectedHead, + build_fully_connect_pytorch_head, + PyTorchSingleCNN, + PyTorchResNet, + PyTorchLoadedModel, +) +from .pytorch.attention import ( + DualSqueezeExciteBlock, + ChannelSqueezeExciteBlock, + SpatialSqueezeExciteBlock, +) +from .pytorch.dataset import PyTorchDataset + + +__all__ = [ + "CTLearnModel", + "SingleCNN", + "ResNet", + "LoadedModel", + "build_fully_connect_keras_head", + "KerasSingleCNN", + "KerasResNet", + "KerasLoadedModel", + "dual_squeeze_excite_block", + "channel_squeeze_excite_block", + "spatial_squeeze_excite_block", + "KerasSequence", + "BasicBlock", + "BottleneckBlock", + "MultiFullyConnectedHead", + "build_fully_connect_pytorch_head", + "PyTorchSingleCNN", + "PyTorchResNet", + "PyTorchLoadedModel", + "DualSqueezeExciteBlock", + "ChannelSqueezeExciteBlock", + "SpatialSqueezeExciteBlock", + "PyTorchDataset", +] \ No newline at end of file diff --git a/ctlearn/core/keras/__init__.py b/ctlearn/core/keras/__init__.py new file mode 100644 index 00000000..5c790f7e --- /dev/null +++ b/ctlearn/core/keras/__init__.py @@ -0,0 +1,27 @@ +""" +ctlearn core Keras functionalities +""" + +from .model import ( + build_fully_connect_keras_head, + KerasSingleCNN, + KerasResNet, + KerasLoadedModel, +) +from .attention import ( + dual_squeeze_excite_block, + channel_squeeze_excite_block, + spatial_squeeze_excite_block, +) +from .sequence import KerasSequence + +__all__ = [ + "build_fully_connect_keras_head", + "KerasSingleCNN", + "KerasResNet", + "KerasLoadedModel", + "dual_squeeze_excite_block", + "channel_squeeze_excite_block", + "spatial_squeeze_excite_block", + "KerasSequence", +] \ No newline at end of file diff --git a/ctlearn/core/attention.py b/ctlearn/core/keras/attention.py similarity index 93% rename from ctlearn/core/attention.py rename to ctlearn/core/keras/attention.py index 0a4d95b3..9069126b 100644 --- a/ctlearn/core/attention.py +++ b/ctlearn/core/keras/attention.py @@ -2,14 +2,14 @@ This module defines the squeeze-excite blocks for channel-wise and/or spatial-wise attention mechanisms. """ -import keras - __all__ = [ "dual_squeeze_excite_block", "channel_squeeze_excite_block", "spatial_squeeze_excite_block", ] +import keras + def dual_squeeze_excite_block(inputs, ratio=16, name=None): """ A channel & spatial (dual) squeeze-excite block. @@ -73,17 +73,13 @@ def channel_squeeze_excite_block(inputs, ratio=4, name=None): Output tensor for the channel squeeze-excite block. """ - # Temp fix for supporting keras2 & keras3 - if int(keras.__version__.split(".")[0]) >= 3: - filters = inputs.shape[-1] - else: - filters = inputs.get_shape().as_list()[-1] + filters = inputs.shape[-1] cse = keras.layers.GlobalAveragePooling2D( keepdims=True, name=name + "_avgpool" )(inputs) cse = keras.layers.Dense( - units=filters // ratio, + units=max(1, filters // ratio), activation="relu", name=name + "_1_dense", )(cse) diff --git a/ctlearn/core/keras/model.py b/ctlearn/core/keras/model.py new file mode 100644 index 00000000..94004f31 --- /dev/null +++ b/ctlearn/core/keras/model.py @@ -0,0 +1,642 @@ +""" +This module defines the ``CTLearnModel`` classes, which holds the basic functionality for creating a Keras model to be used in CTLearn. +""" + +__all__ = [ + "build_fully_connect_keras_head", + "KerasSingleCNN", + "KerasResNet", + "KerasLoadedModel", +] + +import keras + +from ctlearn.core.model import ( + SingleCNN, + ResNet, + LoadedModel, +) +from ctlearn.core.keras.attention import ( + dual_squeeze_excite_block, + channel_squeeze_excite_block, + spatial_squeeze_excite_block, +) + + +def build_fully_connect_keras_head(inputs, layers, activation_function, tasks): + """ + Build the fully connected head for the Keras-based CTLearn model. + + Function to build the fully connected head of the Keras-based CTLearn model using the specified parameters. + + Parameters + ---------- + inputs : keras.layers.Layer + Keras layer of the model. + layers : dict + Dictionary containing the number of neurons (as value) in the fully connected head for each task (as key). + activation_function : dict + Dictionary containing the activation function (as value) for the fully connected head for each task (as key). + tasks : list + List of tasks to build the head for. + + Returns + ------- + logits : dict + Dictionary containing the logits for each task. + """ + logits = {} + for task in tasks: + x = inputs + for i, units in enumerate(layers[task]): + if i != len(layers[task]) - 1: + x = keras.layers.Dense( + units=units, + activation=activation_function[task], + name=f"fc_{task}_{i+1}", + )(x) + else: + x = keras.layers.Dense(units=units, name=task)(x) + logits[task] = keras.layers.Softmax()(x) if task == "type" else x + # 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(tasks) == 1 and tasks[0] == "type": + logits = logits[tasks[0]] + return logits + + +class KerasSingleCNN(SingleCNN): + """ + ``SingleCNN`` is a simple convolutional neural network model. + + This class extends the functionality of ``CTLearnModel`` by implementing + methods to build a simple convolutional neural network model. + """ + + def __init__( + self, + input_shape, + tasks, + config=None, + parent=None, + **kwargs, + ): + super().__init__( + tasks=tasks, + config=config, + parent=parent, + **kwargs, + ) + + # Build the ResNet model backbone + self.backbone_model, self.input_layer = self._build_backbone(input_shape) + backbone_output = self.backbone_model(self.input_layer) + # Build the fully connected head depending on the tasks + self.logits = build_fully_connect_keras_head( + backbone_output, self.head_layers, self.head_activation_function, tasks + ) + # Build the full pipeline model∫ + self.model = keras.Model(self.input_layer, self.logits, name="CTLearn_model") + + def _build_backbone(self, input_shape): + """ + Build the SingleCNN model backbone. + + Function to build the backbone of the SingleCNN model using the specified parameters. + + Parameters + ---------- + input_shape : tuple + Shape of the input data (batch_size, height, width, channels). + + Returns + ------- + backbone_model : keras.Model + Keras model object representing the backbone of the SingleCNN model. + network_input : keras.Input + Keras input layer object for the backbone model. + """ + + # Define the input layer from the input shape + network_input = keras.Input(shape=input_shape) + # Get model arcihtecture parameters for the backbone + filters_list = [layer["filters"] for layer in self.architecture] + kernel_sizes = [layer["kernel_size"] for layer in self.architecture] + numbers_list = [layer["number"] for layer in self.architecture] + + x = network_input + if self.batchnorm: + x = keras.layers.BatchNormalization(momentum=0.99)(x) + + for i, (filters, kernel_size, number) in enumerate( + zip(filters_list, kernel_sizes, numbers_list) + ): + for nr in range(number): + x = keras.layers.Conv2D( + filters=filters, + kernel_size=kernel_size, + padding="same", + activation="relu", + name=f"{self.backbone_name}_conv_{i+1}_{nr+1}", + )(x) + if self.pooling_type is not None: + if self.pooling_type == "max": + x = keras.layers.MaxPool2D( + pool_size=self.pooling_parameters["size"], + strides=self.pooling_parameters["strides"], + name=f"{self.backbone_name}_pool_{i+1}", + )(x) + elif self.pooling_type == "average": + x = keras.layers.AveragePooling2D( + pool_size=self.pooling_parameters["size"], + strides=self.pooling_parameters["strides"], + name=f"{self.backbone_name}_pool_{i+1}", + )(x) + if self.batchnorm: + x = keras.layers.BatchNormalization(momentum=0.99)(x) + + # bottleneck layer + if self.bottleneck_filters is not None: + x = keras.layers.Conv2D( + filters=self.bottleneck_filters, + kernel_size=1, + padding="same", + activation="relu", + name=f"{self.backbone_name}_bottleneck", + )(x) + if self.batchnorm: + x = keras.layers.BatchNormalization(momentum=0.99)(x) + + # Attention mechanism + if self.attention is not None: + if self.attention["mechanism"] == "Dual-SE": + x = dual_squeeze_excite_block( + x, self.attention["reduction_ratio"], name=f"{self.backbone_name}_dse" + ) + elif self.attention["mechanism"] == "Channel-SE": + x = channel_squeeze_excite_block( + x, self.attention["reduction_ratio"], name=f"{self.backbone_name}_cse" + ) + elif self.attention["mechanism"] == "Spatial-SE": + x = spatial_squeeze_excite_block(x, name=f"{self.backbone_name}_sse") + + # Apply global average pooling as the final layer of the backbone + network_output = keras.layers.GlobalAveragePooling2D( + name=self.backbone_name + "_global_avgpool" + )(x) + # Create the backbone model + backbone_model = keras.Model( + network_input, network_output, name=self.backbone_name + ) + return backbone_model, network_input + + +class KerasResNet(ResNet): + """ + ``ResNet`` is a residual neural network model. + + This class extends the functionality of ``CTLearnModel`` by implementing + methods to build a residual neural network model. + """ + + def __init__( + self, + input_shape, + tasks, + config=None, + parent=None, + **kwargs, + ): + super().__init__( + tasks=tasks, + config=config, + parent=parent, + **kwargs, + ) + # Build the ResNet model backbone + self.backbone_model, self.input_layer = self._build_backbone(input_shape) + backbone_output = self.backbone_model(self.input_layer) + # Build the fully connected head depending on the tasks + self.logits = build_fully_connect_keras_head( + backbone_output, self.head_layers, self.head_activation_function, tasks + ) + + self.model = keras.Model(self.input_layer, self.logits, name="CTLearn_model") + + def _build_backbone(self, input_shape): + """ + Build the ResNet model backbone. + + Function to build the backbone of the ResNet model using the specified parameters. + + Parameters + ---------- + input_shape : tuple + Shape of the input data (batch_size, height, width, channels). + + Returns + ------- + backbone_model : keras.Model + Keras model object representing the ResNet backbone. + network_input : keras.Input + Keras input layer object for the backbone model. + """ + # Define the input layer from the input shape + network_input = keras.Input(shape=input_shape) + x = network_input + # Apply initial padding if specified + if self.init_padding > 0: + x = keras.layers.ZeroPadding2D( + padding=self.init_padding, + name=self.backbone_name + "_padding", + )(x) + # Apply initial convolutional layer if specified + if self.init_layer is not None: + x = keras.layers.Conv2D( + filters=self.init_layer["filters"], + kernel_size=self.init_layer["kernel_size"], + strides=self.init_layer["strides"], + name=self.backbone_name + "_conv1_conv", + )(x) + # Apply max pooling if specified + if self.init_max_pool is not None: + x = keras.layers.MaxPool2D( + pool_size=self.init_max_pool["size"], + strides=self.init_max_pool["strides"], + name=self.backbone_name + "_pool1_pool", + )(x) + # Build the residual blocks + engine_output = self._stacked_res_blocks( + x, + architecture=self.architecture, + residual_block_type=self.residual_block_type, + attention=self.attention, + name=self.backbone_name, + ) + # Apply global average pooling as the final layer of the backbone + network_output = keras.layers.GlobalAveragePooling2D( + name=self.backbone_name + "_global_avgpool" + )(engine_output) + # Create the backbone model + backbone_model = keras.Model( + network_input, network_output, name=self.backbone_name + ) + return backbone_model, network_input + + def _stacked_res_blocks( + self, inputs, architecture, residual_block_type, attention, name=None + ): + """ + Build a stack of residual blocks for the CTLearn model. + + This function constructs a stack of residual blocks, which are used to build the backbone of the CTLearn model. + Each residual block consists of a series of convolutional layers with skip connections. + + Parameters + ---------- + inputs : keras.layers.Layer + Input Keras layer to the residual blocks. + architecture : list of dict + List of dictionaries containing the architecture of the ResNet model, which includes: + - Number of filters for the convolutional layers in the residual blocks. + - Number of residual blocks to stack. + residual_block_type : str + Type of residual block to use. Options are 'basic' or 'bottleneck'. + attention : dict + Dictionary containing the configuration parameters for the attention mechanism. + name : str, optional + Label for the model. + + Returns + ------- + x : keras.layers.Layer + Output Keras layer after passing through the stack of residual blocks. + """ + + # Get hyperparameters for the model architecture + filters_list = [layer["filters"] for layer in architecture] + blocks_list = [layer["blocks"] for layer in architecture] + # Build the ResNet model + x = self._stack_fn( + inputs, + filters_list[0], + blocks_list[0], + residual_block_type, + stride=1, + attention=attention, + name=name + "_conv2", + ) + for i, (filters, blocks) in enumerate(zip(filters_list[1:], blocks_list[1:])): + x = self._stack_fn( + x, + filters, + blocks, + residual_block_type, + attention=attention, + name=name + "_conv" + str(i + 3), + ) + return x + + def _stack_fn( + self, + inputs, + filters, + blocks, + residual_block_type, + stride=2, + attention=None, + name=None, + ): + """ + Stack residual blocks for the CTLearn model. + + This function constructs a stack of residual blocks, which are used to build the backbone of the CTLearn model. + Each residual block can be of different types (e.g., basic or bottleneck) and can include attention mechanisms. + + Parameters + ---------- + inputs : keras.layers.Layer + Input tensor to the residual blocks. + filters : int + Number of filters for the bottleneck layer in a block. + blocks : int + Number of residual blocks to stack. + residual_block_type : str + Type of residual block ('basic' or 'bottleneck'). + stride : int, optional + Stride for the first layer in the first block. Default is 2. + attention : dict, optional + Configuration parameters for the attention mechanism. Default is None. + name : str, optional + Label for the stack. Default is None. + + Returns + ------- + keras.layers.Layer + Output tensor for the stacked blocks. + """ + + res_blocks = { + "basic": self._basic_residual_block, + "bottleneck": self._bottleneck_residual_block, + } + + x = res_blocks[residual_block_type]( + inputs, + filters, + stride=stride, + attention=attention, + name=name + "_block1", + ) + for i in range(2, blocks + 1): + x = res_blocks[residual_block_type]( + x, + filters, + conv_shortcut=False, + attention=attention, + name=name + "_block" + str(i), + ) + + return x + + def _basic_residual_block( + self, + inputs, + filters, + kernel_size=3, + stride=1, + conv_shortcut=True, + attention=None, + name=None, + ): + """ + Build a basic residual block for the CTLearn model. + + This function constructs a basic residual block, which is a fundamental building block + of ResNet architectures. The block consists of two convolutional layers with an optional + convolutional shortcut, and can include attention mechanisms. + + Parameters + ---------- + inputs : keras.layers.Layer + Input tensor to the residual block. + filters : int + Number of filters for the convolutional layers. + kernel_size : int, optional + Size of the convolutional kernel. Default is 3. + stride : int, optional + Stride for the convolutional layers. Default is 1. + conv_shortcut : bool, optional + Whether to use a convolutional layer for the shortcut connection. Default is True. + attention : dict, optional + Configuration parameters for the attention mechanism. Default is None. + name : str, optional + Name for the residual block. Default is None. + + Returns + ------- + keras.layers.Layer + Output tensor after applying the residual block. + """ + + if conv_shortcut: + shortcut = keras.layers.Conv2D( + filters=filters, kernel_size=1, strides=stride, name=name + "_0_conv" + )(inputs) + else: + shortcut = inputs + + x = keras.layers.Conv2D( + filters=filters, + kernel_size=kernel_size, + strides=stride, + padding="same", + activation="relu", + name=name + "_1_conv", + )(inputs) + x = keras.layers.Conv2D( + filters=filters, + kernel_size=kernel_size, + padding="same", + activation="relu", + name=name + "_2_conv", + )(x) + + # Attention mechanism + if attention is not None: + if attention["mechanism"] == "Dual-SE": + x = dual_squeeze_excite_block( + x, attention["reduction_ratio"], name=name + "_dse" + ) + elif attention["mechanism"] == "Channel-SE": + x = channel_squeeze_excite_block( + x, attention["reduction_ratio"], name=name + "_cse" + ) + elif attention["mechanism"] == "Spatial-SE": + x = spatial_squeeze_excite_block(x, name=name + "_sse") + + x = keras.layers.Add(name=name + "_add")([x, shortcut]) + x = keras.layers.ReLU(name=name + "_out")(x) + + return x + + def _bottleneck_residual_block( + self, + inputs, + filters, + kernel_size=3, + stride=1, + conv_shortcut=True, + attention=None, + name=None, + ): + """ + Build a bottleneck residual block for the CTLearn model. + + This function constructs a bottleneck residual block, which is a fundamental building block of + ResNet architectures. The block consists of three convolutional layers: a 1x1 convolution to reduce + dimensionality, a 3x3 convolution for main computation, and another 1x1 convolution to restore dimensionality. + It also includes an optional shortcut connection and can include attention mechanisms. + + Parameters + ---------- + inputs : keras.layers.Layer + Input tensor to the residual block. + filters : int + Number of filters for the convolutional layers. + kernel_size : int, optional + Size of the convolutional kernel. Default is 3. + stride : int, optional + Stride for the convolutional layers. Default is 1. + conv_shortcut : bool, optional + Whether to use a convolutional layer for the shortcut connection. Default is True. + attention : dict, optional + Configuration parameters for the attention mechanism. Default is None. + name : str, optional + Name for the residual block. Default is None. + + Returns + ------- + output : keras.layers.Layer + Output layer of the residual block. + """ + + if conv_shortcut: + shortcut = keras.layers.Conv2D( + filters=4 * filters, + kernel_size=1, + strides=stride, + name=name + "_0_conv", + )(inputs) + else: + shortcut = inputs + + x = keras.layers.Conv2D( + filters=filters, + kernel_size=1, + strides=stride, + activation="relu", + name=name + "_1_conv", + )(inputs) + x = keras.layers.Conv2D( + filters=filters, + kernel_size=kernel_size, + padding="same", + activation="relu", + name=name + "_2_conv", + )(x) + x = keras.layers.Conv2D( + filters=4 * filters, kernel_size=1, name=name + "_3_conv" + )(x) + + # Attention mechanism + if attention is not None: + if attention["mechanism"] == "Dual-SE": + x = dual_squeeze_excite_block( + x, attention["reduction_ratio"], name=name + "_dse" + ) + elif attention["mechanism"] == "Channel-SE": + x = channel_squeeze_excite_block( + x, attention["reduction_ratio"], name=name + "_cse" + ) + elif attention["mechanism"] == "Spatial-SE": + x = spatial_squeeze_excite_block(x, name=name + "_sse") + + x = keras.layers.Add(name=name + "_add")([x, shortcut]) + x = keras.layers.ReLU(name=name + "_out")(x) + + return x + + +class KerasLoadedModel(LoadedModel): + """ + ``LoadedModel`` is a pre-trained Keras model. + + This class extends the functionality of ``CTLearnModel`` by implementing + methods to load a pre-trained Keras model. The model can be used as a backbone + for the CTLearn model. + """ + + def __init__( + self, + input_shape, + tasks, + config=None, + parent=None, + **kwargs, + ): + super().__init__( + tasks=tasks, + config=config, + parent=parent, + **kwargs, + ) + + # Load the model from the specified path + self.model = keras.saving.load_model(self.load_model_from) + # Build the ResNet model backbone + self.backbone_model, self.input_layer = self._build_backbone(input_shape) + # Load the fully connected head from the loaded model or build a new one + if self.overwrite_head: + backbone_output = self.backbone_model(self.input_layer) + # Build the fully connected head depending on the tasks + self.logits = build_fully_connect_keras_head( + backbone_output, self.head_layers, self.head_activation_function, tasks + ) + self.model = keras.Model( + self.input_layer, self.logits, name="CTLearn_model" + ) + + def _build_backbone(self, input_shape): + """ + Build the LoadedModel backbone. + + Function to build the backbone of the LoadedModel using the specified parameters. + + Parameters + ---------- + input_shape : tuple + Shape of the input data (batch_size, height, width, channels). + + Returns + ------- + backbone_model : keras.Model + Keras model object representing the LoadedModel backbone. + network_input : keras.Input + Keras input layer object for the backbone model. + """ + + # Define the input layer from the input shape + network_input = keras.Input(shape=input_shape) + # Set the backbone model to be trainable or not + for layer in self.model.layers: + if layer.name.endswith("_block"): + backbone_layer = self.model.get_layer(layer.name) + self.backbone_name = backbone_layer.name + backbone_layer.trainable = self.trainable_backbone + network_output = backbone_layer(network_input) + # Create the backbone model + backbone_model = keras.Model( + network_input, network_output, name=self.backbone_name + ) + return backbone_model, network_input diff --git a/ctlearn/core/loader.py b/ctlearn/core/keras/sequence.py similarity index 99% rename from ctlearn/core/loader.py rename to ctlearn/core/keras/sequence.py index 92fd96c2..3d20a6b0 100644 --- a/ctlearn/core/loader.py +++ b/ctlearn/core/keras/sequence.py @@ -1,12 +1,14 @@ +""" Keras sequence for data loading.""" + +__all__ = ["KerasSequence"] + 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 DLDataLoader(Sequence): +class KerasSequence(Sequence): """ Generates batches for Keras application. diff --git a/ctlearn/core/model.py b/ctlearn/core/model.py index bfac48c2..249cb613 100644 --- a/ctlearn/core/model.py +++ b/ctlearn/core/model.py @@ -2,67 +2,18 @@ This module defines the ``CTLearnModel`` classes, which holds the basic functionality for creating a Keras model to be used in CTLearn. """ -from abc import abstractmethod -import keras - -from ctapipe.core import Component -from ctapipe.core.traits import Bool, Int, CaselessStrEnum, List, Dict, Unicode, Path -from ctlearn.core.attention import ( - dual_squeeze_excite_block, - channel_squeeze_excite_block, - spatial_squeeze_excite_block, -) -from ctlearn.utils import validate_trait_dict - __all__ = [ - "build_fully_connect_head", "CTLearnModel", "SingleCNN", "ResNet", "LoadedModel", ] +from abc import abstractmethod -def build_fully_connect_head(inputs, layers, activation_function, tasks): - """ - Build the fully connected head for the CTLearn model. - - Function to build the fully connected head of the CTLearn model using the specified parameters. - - Parameters - ---------- - inputs : keras.layers.Layer - Keras layer of the model. - layers : dict - Dictionary containing the number of neurons (as value) in the fully connected head for each task (as key). - activation_function : dict - Dictionary containing the activation function (as value) for the fully connected head for each task (as key). - tasks : list - List of tasks to build the head for. - - Returns - ------- - logits : dict - Dictionary containing the logits for each task. - """ - logits = {} - for task in tasks: - x = inputs - for i, units in enumerate(layers[task]): - if i != len(layers[task]) - 1: - x = keras.layers.Dense( - units=units, - activation=activation_function[task], - name=f"fc_{task}_{i+1}", - )(x) - else: - x = keras.layers.Dense(units=units, name=task)(x) - logits[task] = keras.layers.Softmax()(x) if task == "type" else x - # 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(tasks) == 1 and tasks[0] == "type": - logits = logits[tasks[0]] - return logits +from ctapipe.core import Component +from ctapipe.core.traits import Bool, Int, CaselessStrEnum, List, Dict, Unicode, Path +from ctlearn.tools.utils import validate_trait_dict class CTLearnModel(Component): @@ -154,31 +105,31 @@ def __init__( } -@abstractmethod -def _build_backbone(self, input_shape): - """ - Build the backbone of the CTLearn model. + @abstractmethod + def _build_backbone(self, input_shape): + """ + Build the backbone of the CTLearn model. - Function to build the backbone of the CTLearn model using the specified parameters. + Function to build the backbone of the CTLearn model using the specified parameters. - Parameters - ---------- - input_shape : tuple - Shape of the input data (batch_size, height, width, channels). + Parameters + ---------- + input_shape : tuple + Shape of the input data (batch_size, height, width, channels). - Returns - ------- - backbone_model : keras.Model - Keras model object representing the backbone of the CTLearn model. - network_input : keras.Input - Keras input layer object for the backbone model. - """ - pass + Returns + ------- + backbone_model : keras.Model + Keras model object representing the backbone of the CTLearn model. + network_input : keras.Input + Keras input layer object for the backbone model. + """ + pass class SingleCNN(CTLearnModel): """ - ``SingleCNN`` is a simple convolutional neural network model. + ``SingleCNN`` is the base class of a simple convolutional neural network model. This class extends the functionality of ``CTLearnModel`` by implementing methods to build a simple convolutional neural network model. @@ -234,7 +185,6 @@ class SingleCNN(CTLearnModel): def __init__( self, - input_shape, tasks, config=None, parent=None, @@ -251,114 +201,11 @@ def __init__( validate_trait_dict(layer, ["filters", "kernel_size", "number"]) # Validate the pooling parameters trait validate_trait_dict(self.pooling_parameters, ["size", "strides"]) - - # Construct the name of the backbone model by appending "_block" to the model name - self.backbone_name = self.name + "_block" - - # Build the ResNet model backbone - self.backbone_model, self.input_layer = self._build_backbone(input_shape) - backbone_output = self.backbone_model(self.input_layer) # Validate the head trait with the provided tasks validate_trait_dict(self.head_layers, tasks) validate_trait_dict(self.head_activation_function, tasks) - # Build the fully connected head depending on the tasks - self.logits = build_fully_connect_head( - backbone_output, self.head_layers, self.head_activation_function, tasks - ) - - self.model = keras.Model(self.input_layer, self.logits, name="CTLearn_model") - - def _build_backbone(self, input_shape): - """ - Build the SingleCNN model backbone. - - Function to build the backbone of the SingleCNN model using the specified parameters. - - Parameters - ---------- - input_shape : tuple - Shape of the input data (batch_size, height, width, channels). - - Returns - ------- - backbone_model : keras.Model - Keras model object representing the backbone of the SingleCNN model. - network_input : keras.Input - Keras input layer object for the backbone model. - """ - - # Define the input layer from the input shape - network_input = keras.Input(shape=input_shape) - # Get model arcihtecture parameters for the backbone - filters_list = [layer["filters"] for layer in self.architecture] - kernel_sizes = [layer["kernel_size"] for layer in self.architecture] - numbers_list = [layer["number"] for layer in self.architecture] - - x = network_input - if self.batchnorm: - x = keras.layers.BatchNormalization(momentum=0.99)(x) - - for i, (filters, kernel_size, number) in enumerate( - zip(filters_list, kernel_sizes, numbers_list) - ): - for nr in range(number): - x = keras.layers.Conv2D( - filters=filters, - kernel_size=kernel_size, - padding="same", - activation="relu", - name=f"{self.backbone_name}_conv_{i+1}_{nr+1}", - )(x) - if self.pooling_type is not None: - if self.pooling_type == "max": - x = keras.layers.MaxPool2D( - pool_size=self.pooling_parameters["size"], - strides=self.pooling_parameters["strides"], - name=f"{self.backbone_name}_pool_{i+1}", - )(x) - elif self.pooling_type == "average": - x = keras.layers.AveragePooling2D( - pool_size=self.pooling_parameters["size"], - strides=self.pooling_parameters["strides"], - name=f"{self.backbone_name}_pool_{i+1}", - )(x) - if self.batchnorm: - x = keras.layers.BatchNormalization(momentum=0.99)(x) - - # bottleneck layer - if self.bottleneck_filters is not None: - x = keras.layers.Conv2D( - filters=self.bottleneck_filters, - kernel_size=1, - padding="same", - activation="relu", - name=f"{self.backbone_name}_bottleneck", - )(x) - if self.batchnorm: - x = keras.layers.BatchNormalization(momentum=0.99)(x) - - # Attention mechanism - if self.attention is not None: - if self.attention["mechanism"] == "Dual-SE": - x = dual_squeeze_excite_block( - x, self.attention["ratio"], name=f"{self.backbone_name}_dse" - ) - elif self.attention["mechanism"] == "Channel-SE": - x = channel_squeeze_excite_block( - x, self.attention["ratio"], name=f"{self.backbone_name}_cse" - ) - elif self.attention["mechanism"] == "Spatial-SE": - x = spatial_squeeze_excite_block(x, name=f"{self.backbone_name}_sse") - - # Apply global average pooling as the final layer of the backbone - network_output = keras.layers.GlobalAveragePooling2D( - name=self.backbone_name + "_global_avgpool" - )(x) - # Create the backbone model - backbone_model = keras.Model( - network_input, network_output, name=self.backbone_name - ) - return backbone_model, network_input + # Construct the name of the backbone model by appending "_block" to the model name + self.backbone_name = self.name + "_block" class ResNet(CTLearnModel): @@ -416,7 +263,6 @@ class ResNet(CTLearnModel): def __init__( self, - input_shape, tasks, config=None, parent=None, @@ -436,367 +282,11 @@ def __init__( validate_trait_dict(self.init_layer, ["filters", "kernel_size", "strides"]) if self.init_max_pool is not None: validate_trait_dict(self.init_max_pool, ["size", "strides"]) - - # Construct the name of the backbone model by appending "_block" to the model name - self.backbone_name = self.name + "_block" - - # Build the ResNet model backbone - self.backbone_model, self.input_layer = self._build_backbone(input_shape) - backbone_output = self.backbone_model(self.input_layer) # Validate the head traits with the provided tasks validate_trait_dict(self.head_layers, tasks) validate_trait_dict(self.head_activation_function, tasks) - # Build the fully connected head depending on the tasks - self.logits = build_fully_connect_head( - backbone_output, self.head_layers, self.head_activation_function, tasks - ) - - self.model = keras.Model(self.input_layer, self.logits, name="CTLearn_model") - - def _build_backbone(self, input_shape): - """ - Build the ResNet model backbone. - - Function to build the backbone of the ResNet model using the specified parameters. - - Parameters - ---------- - input_shape : tuple - Shape of the input data (batch_size, height, width, channels). - - Returns - ------- - backbone_model : keras.Model - Keras model object representing the ResNet backbone. - network_input : keras.Input - Keras input layer object for the backbone model. - """ - # Define the input layer from the input shape - network_input = keras.Input(shape=input_shape) - # Apply initial padding if specified - if self.init_padding > 0: - network_input = keras.layers.ZeroPadding2D( - padding=self.init_padding, - kernel_size=self.init_layer["kernel_size"], - strides=self.init_layer["strides"], - name=self.backbone_name + "_padding", - )(network_input) - # Apply initial convolutional layer if specified - if self.init_layer is not None: - network_input = keras.layers.Conv2D( - filters=self.init_layer["filters"], - kernel_size=self.init_layer["kernel_size"], - strides=self.init_layer["strides"], - name=self.backbone_name + "_conv1_conv", - )(network_input) - # Apply max pooling if specified - if self.init_max_pool is not None: - network_input = keras.layers.MaxPool2D( - pool_size=self.init_max_pool["size"], - strides=self.init_max_pool["strides"], - name=self.backbone_name + "_pool1_pool", - )(network_input) - # Build the residual blocks - engine_output = self._stacked_res_blocks( - network_input, - architecture=self.architecture, - residual_block_type=self.residual_block_type, - attention=self.attention, - name=self.backbone_name, - ) - # Apply global average pooling as the final layer of the backbone - network_output = keras.layers.GlobalAveragePooling2D( - name=self.backbone_name + "_global_avgpool" - )(engine_output) - # Create the backbone model - backbone_model = keras.Model( - network_input, network_output, name=self.backbone_name - ) - return backbone_model, network_input - - def _stacked_res_blocks( - self, inputs, architecture, residual_block_type, attention, name=None - ): - """ - Build a stack of residual blocks for the CTLearn model. - - This function constructs a stack of residual blocks, which are used to build the backbone of the CTLearn model. - Each residual block consists of a series of convolutional layers with skip connections. - - Parameters - ---------- - inputs : keras.layers.Layer - Input Keras layer to the residual blocks. - architecture : list of dict - List of dictionaries containing the architecture of the ResNet model, which includes: - - Number of filters for the convolutional layers in the residual blocks. - - Number of residual blocks to stack. - residual_block_type : str - Type of residual block to use. Options are 'basic' or 'bottleneck'. - attention : dict - Dictionary containing the configuration parameters for the attention mechanism. - name : str, optional - Label for the model. - - Returns - ------- - x : keras.layers.Layer - Output Keras layer after passing through the stack of residual blocks. - """ - - # Get hyperparameters for the model architecture - filters_list = [layer["filters"] for layer in architecture] - blocks_list = [layer["blocks"] for layer in architecture] - # Build the ResNet model - x = self._stack_fn( - inputs, - filters_list[0], - blocks_list[0], - residual_block_type, - stride=1, - attention=attention, - name=name + "_conv2", - ) - for i, (filters, blocks) in enumerate(zip(filters_list[1:], blocks_list[1:])): - x = self._stack_fn( - x, - filters, - blocks, - residual_block_type, - attention=attention, - name=name + "_conv" + str(i + 3), - ) - return x - - def _stack_fn( - self, - inputs, - filters, - blocks, - residual_block_type, - stride=2, - attention=None, - name=None, - ): - """ - Stack residual blocks for the CTLearn model. - - This function constructs a stack of residual blocks, which are used to build the backbone of the CTLearn model. - Each residual block can be of different types (e.g., basic or bottleneck) and can include attention mechanisms. - - Parameters - ---------- - inputs : keras.layers.Layer - Input tensor to the residual blocks. - filters : int - Number of filters for the bottleneck layer in a block. - blocks : int - Number of residual blocks to stack. - residual_block_type : str - Type of residual block ('basic' or 'bottleneck'). - stride : int, optional - Stride for the first layer in the first block. Default is 2. - attention : dict, optional - Configuration parameters for the attention mechanism. Default is None. - name : str, optional - Label for the stack. Default is None. - - Returns - ------- - keras.layers.Layer - Output tensor for the stacked blocks. - """ - - res_blocks = { - "basic": self._basic_residual_block, - "bottleneck": self._bottleneck_residual_block, - } - - x = res_blocks[residual_block_type]( - inputs, - filters, - stride=stride, - attention=attention, - name=name + "_block1", - ) - for i in range(2, blocks + 1): - x = res_blocks[residual_block_type]( - x, - filters, - conv_shortcut=False, - attention=attention, - name=name + "_block" + str(i), - ) - - return x - - def _basic_residual_block( - self, - inputs, - filters, - kernel_size=3, - stride=1, - conv_shortcut=True, - attention=None, - name=None, - ): - """ - Build a basic residual block for the CTLearn model. - - This function constructs a basic residual block, which is a fundamental building block - of ResNet architectures. The block consists of two convolutional layers with an optional - convolutional shortcut, and can include attention mechanisms. - - Parameters - ---------- - inputs : keras.layers.Layer - Input tensor to the residual block. - filters : int - Number of filters for the convolutional layers. - kernel_size : int, optional - Size of the convolutional kernel. Default is 3. - stride : int, optional - Stride for the convolutional layers. Default is 1. - conv_shortcut : bool, optional - Whether to use a convolutional layer for the shortcut connection. Default is True. - attention : dict, optional - Configuration parameters for the attention mechanism. Default is None. - name : str, optional - Name for the residual block. Default is None. - - Returns - ------- - keras.layers.Layer - Output tensor after applying the residual block. - """ - - if conv_shortcut: - shortcut = keras.layers.Conv2D( - filters=filters, kernel_size=1, strides=stride, name=name + "_0_conv" - )(inputs) - else: - shortcut = inputs - - x = keras.layers.Conv2D( - filters=filters, - kernel_size=kernel_size, - strides=stride, - padding="same", - activation="relu", - name=name + "_1_conv", - )(inputs) - x = keras.layers.Conv2D( - filters=filters, - kernel_size=kernel_size, - padding="same", - activation="relu", - name=name + "_2_conv", - )(x) - - # Attention mechanism - if attention is not None: - if attention["mechanism"] == "Dual-SE": - x = dual_squeeze_excite_block( - x, attention["reduction_ratio"], name=name + "_dse" - ) - elif attention["mechanism"] == "Channel-SE": - x = channel_squeeze_excite_block( - x, attention["reduction_ratio"], name=name + "_cse" - ) - elif attention["mechanism"] == "Spatial-SE": - x = spatial_squeeze_excite_block(x, name=name + "_sse") - - x = keras.layers.Add(name=name + "_add")([shortcut, x]) - x = keras.layers.ReLU(name=name + "_out")(x) - - return x - - def _bottleneck_residual_block( - self, - inputs, - filters, - kernel_size=3, - stride=1, - conv_shortcut=True, - attention=None, - name=None, - ): - """ - Build a bottleneck residual block for the CTLearn model. - - This function constructs a bottleneck residual block, which is a fundamental building block of - ResNet architectures. The block consists of three convolutional layers: a 1x1 convolution to reduce - dimensionality, a 3x3 convolution for main computation, and another 1x1 convolution to restore dimensionality. - It also includes an optional shortcut connection and can include attention mechanisms. - - Parameters - ---------- - inputs : keras.layers.Layer - Input tensor to the residual block. - filters : int - Number of filters for the convolutional layers. - kernel_size : int, optional - Size of the convolutional kernel. Default is 3. - stride : int, optional - Stride for the convolutional layers. Default is 1. - conv_shortcut : bool, optional - Whether to use a convolutional layer for the shortcut connection. Default is True. - attention : dict, optional - Configuration parameters for the attention mechanism. Default is None. - name : str, optional - Name for the residual block. Default is None. - - Returns - ------- - output : keras.layers.Layer - Output layer of the residual block. - """ - - if conv_shortcut: - shortcut = keras.layers.Conv2D( - filters=4 * filters, - kernel_size=1, - strides=stride, - name=name + "_0_conv", - )(inputs) - else: - shortcut = inputs - - x = keras.layers.Conv2D( - filters=filters, - kernel_size=1, - strides=stride, - activation="relu", - name=name + "_1_conv", - )(inputs) - x = keras.layers.Conv2D( - filters=filters, - kernel_size=kernel_size, - padding="same", - activation="relu", - name=name + "_2_conv", - )(x) - x = keras.layers.Conv2D( - filters=4 * filters, kernel_size=1, name=name + "_3_conv" - )(x) - - # Attention mechanism - if attention is not None: - if attention["mechanism"] == "Dual-SE": - x = dual_squeeze_excite_block( - x, attention["reduction_ratio"], name=name + "_dse" - ) - elif attention["mechanism"] == "Channel-SE": - x = channel_squeeze_excite_block( - x, attention["reduction_ratio"], name=name + "_cse" - ) - elif attention["mechanism"] == "Spatial-SE": - x = spatial_squeeze_excite_block(x, name=name + "_sse") - - x = keras.layers.Add(name=name + "_add")([shortcut, x]) - x = keras.layers.ReLU(name=name + "_out")(x) - - return x + # Construct the name of the backbone model by appending "_block" to the model name + self.backbone_name = self.name + "_block" class LoadedModel(CTLearnModel): @@ -831,7 +321,6 @@ class LoadedModel(CTLearnModel): def __init__( self, - input_shape, tasks, config=None, parent=None, @@ -842,54 +331,8 @@ def __init__( parent=parent, **kwargs, ) - - # Load the model from the specified path - self.model = keras.saving.load_model(self.load_model_from) - # Build the ResNet model backbone - self.backbone_model, self.input_layer = self._build_backbone(input_shape) # Load the fully connected head from the loaded model or build a new one if self.overwrite_head: backbone_output = self.backbone_model(self.input_layer) # Validate the head trait with the provided tasks - validate_trait_dict(self.head_layers, tasks) - # Build the fully connected head depending on the tasks - self.logits = build_fully_connect_head( - backbone_output, self.head_layers, self.head_activation_function, tasks - ) - self.model = keras.Model( - self.input_layer, self.logits, name="CTLearn_model" - ) - - def _build_backbone(self, input_shape): - """ - Build the LoadedModel backbone. - - Function to build the backbone of the LoadedModel using the specified parameters. - - Parameters - ---------- - input_shape : tuple - Shape of the input data (batch_size, height, width, channels). - - Returns - ------- - backbone_model : keras.Model - Keras model object representing the LoadedModel backbone. - network_input : keras.Input - Keras input layer object for the backbone model. - """ - - # Define the input layer from the input shape - network_input = keras.Input(shape=input_shape) - # Set the backbone model to be trainable or not - for layer in self.model.layers: - if layer.name.endswith("_block"): - backbone_layer = self.model.get_layer(layer.name) - self.backbone_name = backbone_layer.name - backbone_layer.trainable = self.trainable_backbone - network_output = backbone_layer(network_input) - # Create the backbone model - backbone_model = keras.Model( - network_input, network_output, name=self.backbone_name - ) - return backbone_model, network_input + validate_trait_dict(self.head_layers, tasks) \ No newline at end of file diff --git a/ctlearn/core/pytorch/__init__.py b/ctlearn/core/pytorch/__init__.py new file mode 100644 index 00000000..de657d80 --- /dev/null +++ b/ctlearn/core/pytorch/__init__.py @@ -0,0 +1,34 @@ +""" +ctlearn core PyTorch functionalities +""" + +from .model import ( + BasicBlock, + BottleneckBlock, + MultiFullyConnectedHead, + build_fully_connect_pytorch_head, + PyTorchSingleCNN, + PyTorchResNet, + PyTorchLoadedModel, +) +from .attention import ( + DualSqueezeExciteBlock, + ChannelSqueezeExciteBlock, + SpatialSqueezeExciteBlock, +) +from .dataset import PyTorchDataset + + +__all__ = [ + "BasicBlock", + "BottleneckBlock", + "MultiFullyConnectedHead", + "build_fully_connect_pytorch_head", + "PyTorchSingleCNN", + "PyTorchResNet", + "PyTorchLoadedModel", + "DualSqueezeExciteBlock", + "ChannelSqueezeExciteBlock", + "SpatialSqueezeExciteBlock", + "PyTorchDataset", +] \ No newline at end of file diff --git a/ctlearn/core/pytorch/attention.py b/ctlearn/core/pytorch/attention.py new file mode 100644 index 00000000..1e1c5cd3 --- /dev/null +++ b/ctlearn/core/pytorch/attention.py @@ -0,0 +1,69 @@ +""" +This module defines the squeeze-excite blocks for channel-wise and/or spatial-wise attention mechanisms in PyTorch. +""" + +__all__ = [ + "DualSqueezeExciteBlock", + "ChannelSqueezeExciteBlock", + "SpatialSqueezeExciteBlock", +] + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class DualSqueezeExciteBlock(nn.Module): + """ + A channel & spatial (dual) squeeze-excite block in PyTorch. + Concurrently applies channel and spatial scaling, then sums the results. + """ + def __init__(self, in_channels, ratio=16): + super().__init__() + self.cse = ChannelSqueezeExciteBlock(in_channels=in_channels, ratio=ratio) + self.sse = SpatialSqueezeExciteBlock(in_channels=in_channels) + + def forward(self, x): + # Combines cse and sse by element-wise addition + return self.cse(x) + self.sse(x) + + +class ChannelSqueezeExciteBlock(nn.Module): + """ + A channel-wise squeeze-excite (cSE) block in PyTorch. + """ + def __init__(self, in_channels, ratio=4): + super().__init__() + reduced_channels = max(1, in_channels // ratio) + # Using nn.Linear to match Keras Dense layers + self.fc1 = nn.Linear(in_channels, reduced_channels, bias=True) + self.fc2 = nn.Linear(reduced_channels, in_channels, bias=True) + + def forward(self, x): + batch_size, channels, _, _ = x.shape + # Global Average Pooling keeping dimensions: (B, C, H, W) -> (B, C, 1, 1) + squeeze = F.adaptive_avg_pool2d(x, (1, 1)) + # Flatten for Linear layers: (B, C, 1, 1) -> (B, C) + squeeze = squeeze.view(batch_size, channels) + # Dense projections with ReLU and Sigmoid + excitation = F.relu(self.fc1(squeeze)) + excitation = torch.sigmoid(self.fc2(excitation)) + # Reshape back to broadcast across spatial dimensions: (B, C) -> (B, C, 1, 1) + excitation = excitation.view(batch_size, channels, 1, 1) + # Scale input tensor + return x * excitation + +class SpatialSqueezeExciteBlock(nn.Module): + """ + A spatial squeeze-excite (sSE) block in PyTorch. + """ + def __init__(self, in_channels): + super().__init__() + # A 1x1 convolution projecting channels down to 1 spatial mask + self.spatial_conv = nn.Conv2d(in_channels, 1, kernel_size=1, bias=True) + + def forward(self, x): + # Create a spatial landscape mask via sigmoid + spatial_mask = torch.sigmoid(self.spatial_conv(x)) + # Multiply input tensor element-wise across spatial layout + return x * spatial_mask \ No newline at end of file diff --git a/ctlearn/core/pytorch/dataset.py b/ctlearn/core/pytorch/dataset.py new file mode 100644 index 00000000..f2fa7f22 --- /dev/null +++ b/ctlearn/core/pytorch/dataset.py @@ -0,0 +1,209 @@ +""" PyTorch dataset for data loading.""" + + +__all__ = ["PyTorchDataset"] + +import numpy as np +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset + +from dl1_data_handler.reader import ProcessType + + +class PyTorchDataset(Dataset): + """ + Generates items/batches for PyTorch application based on DLDataReader. + """ + + def __init__( + self, + DLDataReader, + indices, + tasks, + sort_by_intensity=False, + stack_telescope_images=False, + ): + super().__init__() + self.DLDataReader = DLDataReader + self.indices = list(indices) + self.tasks = tasks + self.sort_by_intensity = sort_by_intensity + self.stack_telescope_images = stack_telescope_images + + # Convert Keras (H, W, C) input_shape to PyTorch (C, H, W) + if self.DLDataReader.__class__.__name__ != "DLFeatureVectorReader": + # 2. Extract base Keras shape safely (H, W, C) + if self.DLDataReader.mode == "mono": + keras_shape = self.DLDataReader.input_shape + elif self.DLDataReader.mode == "stereo": + first_tel = next(iter(self.DLDataReader.selected_telescopes)) + keras_shape = self.DLDataReader.input_shape[first_tel] + # In Keras, 4D input shapes are usually (batch/num_tels, H, W, C) + if self.stack_telescope_images: + num_tels, h, w, c = keras_shape + keras_shape = (h, w, num_tels * c) + + # Permute Keras (H, W, C) to PyTorch (C, H, W) via unpacking + h, w, c = keras_shape + self.input_shape = (c, h, w) + + def __len__(self): + return len(self.indices) + + def __getitem__(self, idx): + """ + Retrieves a single data item at the given index. + Note: If passed a slice/list of indices via custom batching, + it falls back to loading as a mini-batch. + """ + # Support both single index lookup and batch slice lookup + if isinstance(idx, (int, np.integer)): + batch_indices = [self.indices[idx]] + else: + batch_indices = [self.indices[i] for i in idx] + + 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) + + # Remove explicit batch dim if caller requested a single index + if isinstance(idx, (int, np.integer)): + if isinstance(features, torch.Tensor): + features = features.squeeze(0) + if isinstance(labels, dict): + labels = {k: v.squeeze(0) for k, v in labels.items()} + elif isinstance(labels, torch.Tensor): + labels = labels.squeeze(0) + + return features, labels + + def _get_mono_item(self, batch): + labels = {} + # Transpose raw batch: (B, H, W, C) -> (B, C, H, W) + raw_features = torch.from_numpy(batch["features"].data).float() + features = raw_features.permute(0, 3, 1, 2) + + # Construct task tensors + if "type" in self.tasks: + # Return class indices directly as 1D long tensor + type_labels = torch.from_numpy(batch["true_shower_primary_class"].data).long() + labels["type"] = type_labels + if len(self.tasks) == 1: + labels = type_labels + + if "energy" in self.tasks: + energy_tensor = torch.from_numpy(batch["log_true_energy"].data).float() + if isinstance(labels, dict): + labels["energy"] = energy_tensor + else: + labels = energy_tensor + + if "skydirection" in self.tasks: + sky = np.stack((batch["fov_lon"].data, batch["fov_lat"].data), axis=1) + sky_tensor = torch.from_numpy(sky).float() + if isinstance(labels, dict): + labels["skydirection"] = sky_tensor + + if "cameradirection" in self.tasks: + cam = np.stack( + (batch["cam_coord_offset_x"].data, batch["cam_coord_offset_y"].data), + axis=1, + ) + cam_tensor = torch.from_numpy(cam).float() + if isinstance(labels, dict): + labels["cameradirection"] = cam_tensor + + return features, labels + + def _get_stereo_item(self, batch): + 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 = [], [] + cam_coord_offset_x, cam_coord_offset_y = [], [] + + for group_element in batch_grouped.groups: + if "features" in batch.colnames: + if self.sort_by_intensity: + group_element.sort(["hillas_intensity"], reverse=True) + if self.stack_telescope_images: + plain_features = group_element["features"].data + stacked_features = np.concatenate( + [plain_features[i] for i in range(plain_features.shape[0])], + axis=-1, + ) + features.append(stacked_features) + else: + features.append(group_element["features"].data) + + 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 + ) + + 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) + + # Construct task tensors + if "type" in self.tasks: + # 1D Tensor of class indices: shape (batch_size,) + type_labels = torch.tensor(true_shower_primary_class, dtype=torch.long) + labels["type"] = type_labels + if len(self.tasks) == 1: + labels = type_labels + + if "energy" in self.tasks: + energy_tensor = torch.tensor(log_true_energy, dtype=torch.float32) + if isinstance(labels, dict): + labels["energy"] = energy_tensor + + if "skydirection" in self.tasks: + sky = np.stack((np.array(fov_lon), np.array(fov_lat)), axis=1) + sky_tensor = torch.from_numpy(sky).float() + if isinstance(labels, dict): + labels["skydirection"] = sky_tensor + + if "cameradirection" in self.tasks: + cam = np.stack( + (np.array(cam_coord_offset_x), np.array(cam_coord_offset_y)), + axis=1, + ) + cam_tensor = torch.from_numpy(cam).float() + if isinstance(labels, dict): + labels["cameradirection"] = cam_tensor + + # Permute feature axes to PyTorch channel conventions + if "features" in batch.colnames: + raw_features = torch.tensor(np.array(features), dtype=torch.float32) + # Shape: (B, H, W, C) -> Permute to (B, C, H, W) + features = raw_features.permute(0, 3, 1, 2) + if "mono_feature_vectors" in batch.colnames: + features = torch.tensor(np.array(mono_feature_vectors), dtype=torch.float32) + if "stereo_feature_vectors" in batch.colnames: + features = torch.tensor(np.array(stereo_feature_vectors), dtype=torch.float32) + + return features, labels \ No newline at end of file diff --git a/ctlearn/core/pytorch/model.py b/ctlearn/core/pytorch/model.py new file mode 100644 index 00000000..589895b1 --- /dev/null +++ b/ctlearn/core/pytorch/model.py @@ -0,0 +1,520 @@ +""" +This module defines the ``CTLearnModel`` classes, which holds the basic functionality for creating a PyTorch model to be used in CTLearn. +""" + +__all__ = [ + "BasicBlock", + "BottleneckBlock", + "MultiFullyConnectedHead", + "build_fully_connect_pytorch_head", + "PyTorchSingleCNN", + "PyTorchResNet", + "PyTorchLoadedModel", +] + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# Assuming these custom attention blocks are updated to return torch.nn.Module or used dynamically +from ctlearn.core.model import ( + SingleCNN, + ResNet, + LoadedModel, +) +from ctlearn.core.pytorch.attention import ( + DualSqueezeExciteBlock, + ChannelSqueezeExciteBlock, + SpatialSqueezeExciteBlock, +) + + +class MultiFullyConnectedHead(nn.Module): + """ + A PyTorch container module to hold the multi-task fully connected heads. + """ + + def __init__(self, heads_dict, single_output_task=None): + super().__init__() + # Save the dict with tasks info in an attribute + self.heads_dict = heads_dict + # Sanitize keys because 'type' conflicts with nn.Module.type() method + self._task_mapping = { + task: f"head_{task}" if hasattr(nn.Module, task) else task + for task in self.heads_dict.keys() + } + sanitized_heads = { + self._task_mapping[task]: module for task, module in heads_dict.items() + } + self.heads = nn.ModuleDict(sanitized_heads) + self.single_output_task = single_output_task + + def forward(self, x): + # Flatten backbone output if spatially aggregated (B, C, 1, 1) -> (B, C) + if x.dim() > 2: + x = torch.flatten(x, start_dim=1) + + outputs = {} + + for original_task, internal_key in self._task_mapping.items(): + head = self.heads[internal_key] + out = head(x) + + if original_task == "type": + outputs[original_task] = F.softmax(out, dim=-1) + else: + outputs[original_task] = out + + # If only a single task is present, return a single output tensor or dictionary + if self.single_output_task: + return outputs[self.single_output_task] + + return outputs + + +def build_fully_connect_pytorch_head(in_features, layers, activation_function, tasks): + """ + Build the fully connected head for the PyTorch-based CTLearn model. + """ + heads = {} + + # Activation mapping from Keras string to PyTorch Module + act_map = { + "relu": nn.ReLU, + "tanh": nn.Tanh, + "sigmoid": nn.Sigmoid + } + + for task in tasks: + task_layers = [] + current_features = in_features + + for i, units in enumerate(layers[task]): + task_layers.append(nn.Linear(current_features, units)) + if i != len(layers[task]) - 1: + act_cls = act_map.get(activation_function[task].lower(), nn.ReLU) + task_layers.append(act_cls()) + current_features = units + + heads[task] = nn.Sequential(*task_layers) + + single_output_task = tasks[0] if (len(tasks) == 1 and tasks[0] == "type") else None + return MultiFullyConnectedHead(heads, single_output_task=single_output_task) + + +class FullModelPipeline(nn.Module): + """ + Combines the backbone and multi-task heads into a unified executable nn.Module pipeline. + """ + def __init__(self, backbone, head): + super().__init__() + self.backbone = backbone + self.head = head + + def forward(self, x): + features = self.backbone(x) + return self.head(features), features + + +class PyTorchSingleCNN(SingleCNN): + """ + ``SingleCNN`` is a simple convolutional neural network model implemented in PyTorch. + """ + + def __init__(self, input_shape, tasks, config=None, parent=None, **kwargs): + super().__init__(tasks=tasks, config=config, parent=parent, **kwargs) + + # Build modules + self.backbone_model, out_features = self._build_backbone(input_shape) + self.logits_head = build_fully_connect_pytorch_head( + out_features, self.head_layers, self.head_activation_function, tasks + ) + # Final native PyTorch module pipeline saved into self.model + self.model = FullModelPipeline(self.backbone_model, self.logits_head) + + def _build_backbone(self, input_shape): + # input_shape format: (channels, height, width) + in_channels = input_shape[0] + modules = [] + + if self.batchnorm: + modules.append(nn.BatchNorm2d(in_channels, momentum=0.01, eps=1e-3)) # PyTorch momentum = 1 - Keras momentum + + for i, layer in enumerate(self.architecture): + filters = layer["filters"] + kernel_size = layer["kernel_size"] + number = layer["number"] + + for nr in range(number): + # padding="same" calculates padding dynamically in PyTorch based on kernel size + padding = kernel_size // 2 + modules.append(nn.Conv2d(in_channels, filters, kernel_size=kernel_size, padding=padding)) + modules.append(nn.ReLU()) + in_channels = filters + + if self.pooling_type is not None: + p_size = self.pooling_parameters["size"] + p_stride = self.pooling_parameters["strides"] + if self.pooling_type == "max": + modules.append(nn.MaxPool2d(kernel_size=p_size, stride=p_stride)) + elif self.pooling_type == "average": + modules.append(nn.AvgPool2d(kernel_size=p_size, stride=p_stride)) + + if self.batchnorm: + modules.append(nn.BatchNorm2d(in_channels, momentum=0.01, eps=1e-3)) + + if self.bottleneck_filters is not None: + modules.append(nn.Conv2d(in_channels, self.bottleneck_filters, kernel_size=1)) + modules.append(nn.ReLU()) + in_channels = self.bottleneck_filters + if self.batchnorm: + modules.append(nn.BatchNorm2d(in_channels, momentum=0.01, eps=1e-3)) + + if self.attention is not None: + mech = self.attention.get("mechanism") + ratio = self.attention.get("reduction_ratio", 16) + if mech == "Dual-SE": + attention_layer = DualSqueezeExciteBlock(in_channels=in_channels, ratio=ratio) + elif mech == "Channel-SE": + attention_layer = ChannelSqueezeExciteBlock(in_channels=in_channels, ratio=ratio) + elif mech == "Spatial-SE": + attention_layer = SpatialSqueezeExciteBlock(in_channels=in_channels) + modules.append(attention_layer) + + # Perform global 2D average pooling + modules.append(nn.AdaptiveAvgPool2d((1, 1))) + modules.append(nn.Flatten(start_dim=1)) + + return nn.Sequential(*modules), in_channels + + +class BasicBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride=1, conv_shortcut=True, attention=None): + super().__init__() + self.conv_shortcut = conv_shortcut + self.attention_config = attention + # Projection Shortcut Branch + if conv_shortcut: + self.shortcut = nn.Conv2d( + in_channels, out_channels, kernel_size=1, stride=stride, bias=True + ) + else: + self.shortcut = None + # Main Branch Convolutions (Matching Keras _1_conv and _2_conv) + self.conv1 = nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=True + ) + self.conv2 = nn.Conv2d( + out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True + ) + # Setup the attention mechanism + self.setup_attention(out_channels) + + def setup_attention(self, channels): + self.attn_layer = None + if self.attention_config: + mech = self.attention_config.get("mechanism") + ratio = self.attention_config.get("reduction_ratio", 16) + if mech == "Dual-SE": + self.attn_layer = DualSqueezeExciteBlock(in_channels=channels, ratio=ratio) + elif mech == "Channel-SE": + self.attn_layer = ChannelSqueezeExciteBlock(in_channels=channels, ratio=ratio) + elif mech == "Spatial-SE": + self.attn_layer = SpatialSqueezeExciteBlock(in_channels=channels) + + def forward(self, x): + # Shortcut path + if self.conv_shortcut and self.shortcut is not None: + identity = self.shortcut(x) + else: + identity = x + + # Main path (Matches Keras activation order: conv1 -> relu -> conv2) + out = F.relu(self.conv1(x)) + out = self.conv2(out) + + if self.attn_layer is not None: + out = self.attn_layer(out) + + out += identity + return F.relu(out) + + +class BottleneckBlock(nn.Module): + def __init__(self, in_channels, base_filters, stride=1, conv_shortcut=True, attention=None): + super().__init__() + self.conv_shortcut = conv_shortcut + self.attention_config = attention + + # Shortcut connection + if conv_shortcut: + self.shortcut = nn.Conv2d( + in_channels, 4 * base_filters, kernel_size=1, stride=stride, bias=False + ) + else: + self.shortcut = nn.Identity() + # Main branch convolutions matching Keras layout: + self.conv1 = nn.Conv2d( + in_channels, base_filters, kernel_size=1, stride=stride, bias=False + ) + # Keras _2_conv is 3x3 spatial convolution with stride=1 (since stride is handled in _1_conv) + self.conv2 = nn.Conv2d( + base_filters, base_filters, kernel_size=3, stride=1, padding=1, bias=False + ) + # Keras _3_conv restores channels back to 4 * base_filters + self.conv3 = nn.Conv2d( + base_filters, 4 * base_filters, kernel_size=1, bias=False + ) + # Setup the attention mechanism + self.setup_attention(4 * base_filters) + + def setup_attention(self, channels): + self.attn_layer = None + if self.attention_config: + mech = self.attention_config["mechanism"] + ratio = self.attention_config.get("reduction_ratio", 16) + if mech == "Dual-SE": + self.attn_layer = DualSqueezeExciteBlock(in_channels=channels, ratio=ratio) + elif mech == "Channel-SE": + self.attn_layer = ChannelSqueezeExciteBlock(in_channels=channels, ratio=ratio) + elif mech == "Spatial-SE": + self.attn_layer = SpatialSqueezeExciteBlock(in_channels=channels) + + def forward(self, x): + identity = self.shortcut(x) + + # Matches Keras sequence: conv1 (with stride) -> relu -> conv2 -> relu -> conv3 + out = F.relu(self.conv1(x)) + out = F.relu(self.conv2(out)) + out = self.conv3(out) + + if self.attn_layer: + out = self.attn_layer(out) + + out += identity + return F.relu(out) + +class PyTorchResNet(ResNet): + """ + ``PyTorchResNet`` is a residual neural network model implemented in PyTorch. + """ + + def __init__(self, input_shape, tasks, config=None, parent=None, **kwargs): + super().__init__(tasks=tasks, config=config, parent=parent, **kwargs) + + # Build PyTorch backbone and track final out_features channel size + self.backbone_model, out_features = self._build_backbone(input_shape) + # Build the fully connected head + self.logits_head = build_fully_connect_pytorch_head( + out_features, self.head_layers, self.head_activation_function, tasks + ) + # Unify into our structural pipeline wrapper module + self.model = FullModelPipeline(self.backbone_model, self.logits_head) + + def _build_backbone(self, input_shape): + in_channels = input_shape[0] + modules = [] + + # Initial Zero Padding + if getattr(self, "init_padding", 0) > 0: + modules.append(nn.ZeroPad2d(self.init_padding)) + + # Initial Conv Layer + if getattr(self, "init_layer", None) is not None: + out_ch = self.init_layer["filters"] + k_size = self.init_layer["kernel_size"] + stride = self.init_layer["strides"] + padding = k_size // 2 + + modules.append( + nn.Conv2d( + in_channels, + out_ch, + kernel_size=k_size, + stride=stride, + padding=padding, + bias=True, # Match Keras default bias if applicable + ) + ) + modules.append(nn.ReLU()) + in_channels = out_ch + + # Initial Max Pooling + if getattr(self, "init_max_pool", None) is not None: + p_size = self.init_max_pool["size"] + p_stride = self.init_max_pool["strides"] + modules.append( + nn.MaxPool2d(kernel_size=p_size, stride=p_stride, padding=0) + ) + + # Assemble Stacked Residual Architecture blocks + res_blocks, final_channels = self._stacked_res_blocks( + in_channels, + architecture=self.architecture, + residual_block_type=self.residual_block_type, + attention=self.attention, + ) + modules.extend(res_blocks) + + # Perform global 2D average pooling + modules.append(nn.AdaptiveAvgPool2d((1, 1))) + modules.append(nn.Flatten(start_dim=1)) + + return nn.Sequential(*modules), final_channels + + def _stacked_res_blocks(self, in_channels, architecture, residual_block_type, attention): + blocks_list = [] + current_channels = in_channels + + filters_list = [layer["filters"] for layer in architecture] + blocks_count = [layer["blocks"] for layer in architecture] + + # First layer block sequence (stride=1) + blocks_list.extend( + self._stack_fn( + current_channels, + filters_list[0], + blocks_count[0], + residual_block_type, + stride=1, + attention=attention, + ) + ) + + multiplier = 4 if residual_block_type == "bottleneck" else 1 + current_channels = filters_list[0] * multiplier + + # Subsequent downsampling levels (stride=2) + for filters, blocks in zip(filters_list[1:], blocks_count[1:]): + blocks_list.extend( + self._stack_fn( + current_channels, + filters, + blocks, + residual_block_type, + stride=2, + attention=attention, + ) + ) + current_channels = filters * multiplier + + return blocks_list, current_channels + + def _stack_fn(self, in_channels, filters, blocks, residual_block_type, stride=2, attention=None): + stack = [] + # Bottleneck blocks expand channels by 4x; Basic blocks do not expand. + multiplier = 4 if residual_block_type == "bottleneck" else 1 + out_channels = filters * multiplier + + def build_block(in_c, s): + # Only use a conv shortcut if channels change or if downsampling (stride > 1) + needs_shortcut = (in_c != out_channels) or (s != 1) + if residual_block_type == "basic": + return BasicBlock( + in_channels=in_c, + out_channels=filters, + stride=s, + conv_shortcut=needs_shortcut, + attention=attention, + ) + else: + return BottleneckBlock( + in_channels=in_c, + base_filters=filters, + stride=s, + conv_shortcut=needs_shortcut, + attention=attention, + ) + + # First block transition + stack.append(build_block(in_channels, s=stride)) + + # Remaining blocks in the layer + for _ in range(1, blocks): + stack.append(build_block(out_channels, s=1)) + + return stack + + +class PyTorchLoadedModel(LoadedModel): + """ + ``PyTorchLoadedModel`` is a pre-trained PyTorch model wrapper. + + This class loads a pre-trained PyTorch ``nn.Module`` object directly from disk + and uses its backbone and feature representation. + """ + + def __init__( + self, + input_shape, + tasks, + config=None, + parent=None, + **kwargs, + ): + super().__init__( + tasks=tasks, + config=config, + parent=parent, + **kwargs, + ) + + # 1. Load object directly from disk + loaded_object = torch.load(self.load_model_from, map_location="cpu", weights_only=False) + + # 2. Strict validation: ensure the loaded object is an nn.Module + if not isinstance(loaded_object, nn.Module): + raise TypeError( + f"Expected a PyTorch 'nn.Module' object at '{self.load_model_from}', " + f"but got '{type(loaded_object).__name__}'. " + "Ensure the model was saved via 'torch.save(model, path)' rather than 'torch.save(model.state_dict(), path)'." + ) + + self.loaded_model = loaded_object + + # 3. Extract backbone and feature depth + self.backbone_model, out_features = self._build_backbone(input_shape) + + # 4. Construct output pipeline + if self.overwrite_head: + self.logits_head = build_fully_connect_pytorch_head( + out_features, self.head_layers, self.head_activation_function, tasks + ) + self.model = FullModelPipeline(self.backbone_model, self.logits_head) + else: + self.model = self.loaded_model + + def _build_backbone(self, input_shape): + """ + Extract the backbone from the loaded nn.Module and set parameter trainability. + + Parameters + ---------- + input_shape : tuple + Shape of the input data (channels, height, width). + + Returns + ------- + backbone_model : nn.Module + PyTorch module representing the backbone. + out_features : int + Number of output feature channels from the backbone. + """ + # Extract backbone attribute if present, otherwise use full model + if hasattr(self.loaded_model, "backbone"): + backbone_model = self.loaded_model.backbone + else: + backbone_model = self.loaded_model + + # Configure parameter trainability + for param in backbone_model.parameters(): + param.requires_grad = self.trainable_backbone + + # Extract output dimensions for head construction + if hasattr(self.loaded_model, "head") and hasattr(self.loaded_model.head, "heads"): + first_head = next(iter(self.loaded_model.head.heads.values())) + out_features = first_head[0].in_features + else: + out_features = self.head_layers[self.tasks[0]][0] + + return backbone_model, out_features \ No newline at end of file diff --git a/ctlearn/core/tests/test_attention.py b/ctlearn/core/tests/test_attention.py new file mode 100644 index 00000000..f519943c --- /dev/null +++ b/ctlearn/core/tests/test_attention.py @@ -0,0 +1,59 @@ +import pytest +import torch +import keras +import tensorflow as tf + +from ctlearn.core.keras.model import ( + channel_squeeze_excite_block, + spatial_squeeze_excite_block, + dual_squeeze_excite_block, +) +from ctlearn.core.pytorch.attention import( + ChannelSqueezeExciteBlock, + SpatialSqueezeExciteBlock, + DualSqueezeExciteBlock, +) + +def build_keras_se_model(block_fn, input_shape, **kwargs): + """Wraps a Keras functional squeeze-excite block in a Keras Model.""" + inputs = keras.Input(shape=input_shape) + outputs = block_fn(inputs, name="se_block", **kwargs) + return keras.Model(inputs=inputs, outputs=outputs) + + +@pytest.mark.parametrize( + "k_fn, p_class, kwargs", + [ + (channel_squeeze_excite_block, ChannelSqueezeExciteBlock, {"ratio": 4}), + (spatial_squeeze_excite_block, SpatialSqueezeExciteBlock, {}), + (dual_squeeze_excite_block, DualSqueezeExciteBlock, {"ratio": 16}), + ], +) +@pytest.mark.parametrize( + "batch, height, width, channels", + [ + (1, 8, 8, 16), + (4, 32, 32, 64), + ], +) +def test_output_shape_parity(k_fn, p_class, kwargs, batch, height, width, channels): + """Verifies that output shapes match between Keras (BHWC) and PyTorch (BCHW).""" + # Keras Input: (Batch, H, W, C) + x_k = tf.random.normal((batch, height, width, channels)) + k_model = build_keras_se_model(k_fn, input_shape=(height, width, channels), **kwargs) + k_out = k_model(x_k) + + # PyTorch Input: (Batch, C, H, W) + x_p = torch.randn(batch, channels, height, width) + p_module = p_class(in_channels=channels, **kwargs) + p_module.eval() + with torch.no_grad(): + p_out = p_module(x_p) + + # Check output shape correspondence + assert k_out.shape == (batch, height, width, channels) + assert p_out.shape == (batch, channels, height, width) + + # Verify equivalent dimensions (transpose PyTorch output to BHWC) + p_out_bhwc = p_out.permute(0, 2, 3, 1) + assert k_out.shape == p_out_bhwc.shape diff --git a/ctlearn/core/tests/test_loader.py b/ctlearn/core/tests/test_loader.py deleted file mode 100644 index 7fe71900..00000000 --- a/ctlearn/core/tests/test_loader.py +++ /dev/null @@ -1,37 +0,0 @@ -from traitlets.config.loader import Config - -from dl1_data_handler.reader import DLImageReader -from ctlearn.core.loader import DLDataLoader - - -def test_data_loader(dl1_gamma_file): - """check""" - # Create a configuration suitable for the test - config = Config( - { - "DLImageReader": { - "allowed_tels": [4], - "focal_length_choice": "EQUIVALENT", - }, - } - ) - # Create an image reader - dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) - # Create a data loader - dl1_loader = DLDataLoader( - DLDataReader=dl1_reader, - indices=[0], - tasks=["type", "energy", "cameradirection", "skydirection"], - batch_size=1, - ) - # Get the features and labels fgrom the data loader for one batch - features, labels = dl1_loader[0] - # Check that all the correct labels are present - assert ( - "type" in labels - and "energy" in labels - and "cameradirection" in labels - and "skydirection" in labels - ) - # Check the shape of the features - assert features.shape == (1, 110, 110, 2) diff --git a/ctlearn/core/tests/test_loaders.py b/ctlearn/core/tests/test_loaders.py new file mode 100644 index 00000000..94590829 --- /dev/null +++ b/ctlearn/core/tests/test_loaders.py @@ -0,0 +1,54 @@ +import pytest +from torch.utils.data import DataLoader +from traitlets.config.loader import Config + +from dl1_data_handler.reader import DLImageReader +from ctlearn.core.keras.sequence import KerasSequence +from ctlearn.core.pytorch.dataset import PyTorchDataset + + +@pytest.mark.parametrize( + ("dataloader_cls", "expected_features_shape"), + [ + (KerasSequence, (1, 110, 110, 2)), + (PyTorchDataset, (1, 2, 110, 110)), + ], + ids=["Keras", "PyTorch"], +) +def test_data_loading(dl1_gamma_file, dataloader_cls, expected_features_shape): + """Check the data loading for the Keras and PyTorch frameworks""" + # Create a configuration suitable for the test + config = Config( + { + "DLImageReader": { + "allowed_tels": [4], + "focal_length_choice": "EQUIVALENT", + }, + } + ) + # Create an image reader + dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) + # Initialize the PyTorch Dataset or the Keras Sequence + dl1_dataset = dataloader_cls( + DLDataReader=dl1_reader, + indices=[0], + tasks=["type", "energy", "cameradirection", "skydirection"], + ) + # For PyTorch, wrap in DataLoader to apply batch dimension (dim=0) + if issubclass(dataloader_cls, PyTorchDataset): + batch_loader = DataLoader(dl1_dataset, batch_size=1) + features, labels = next(iter(batch_loader)) + # Optional: convert PyTorch tensor to numpy if labels/features are Tensors + if hasattr(features, "numpy"): + features = features.numpy() + else: + features, labels = dl1_dataset[0] + # Check that all the correct labels are present + assert ( + "type" in labels + and "energy" in labels + and "cameradirection" in labels + and "skydirection" in labels + ) + # Check the shape of the features match the expected ones + assert features.shape == expected_features_shape diff --git a/ctlearn/core/tests/test_models.py b/ctlearn/core/tests/test_models.py new file mode 100644 index 00000000..68c77687 --- /dev/null +++ b/ctlearn/core/tests/test_models.py @@ -0,0 +1,223 @@ +import pytest +import numpy as np +import keras +import torch +import torch.nn as nn + +from ctlearn.core.keras.model import KerasResNet, KerasSingleCNN +from ctlearn.core.pytorch.model import PyTorchResNet, PyTorchSingleCNN + +rng = np.random.default_rng(42) + +@pytest.fixture +def common_config(): + return { + "tasks": ["type", "energy", "cameradirection"], + "input_shape_keras": (110, 110, 2), # (H, W, C) + "input_shape_pytorch": (2, 110, 110), # (C, H, W) + "kwargs": { + "head_layers": { + "type": [64, 2], + "energy": [64, 1], + "cameradirection": [64, 2], + }, + "head_activation_function": { + "type": "relu", + "energy": "relu", + "cameradirection": "tanh", + }, + }, + } + +@pytest.mark.parametrize("batchnorm", [True, False]) +@pytest.mark.parametrize( + "attention", + [ + {"mechanism": None}, + {"mechanism": "Channel-SE", "reduction_ratio": 8}, + {"mechanism": "Spatial-SE"}, + {"mechanism": "Dual-SE", "reduction_ratio": 32}, + ], +) +def test_SingleCNN_model_structure_parity(common_config, batchnorm, attention): + """Verify that Keras and PyTorch models have matching layer counts and weight shapes.""" + tasks = common_config["tasks"] + kwargs = common_config["kwargs"].copy() + kwargs["architecture"] = [ + {"filters": 32, "kernel_size": 2, "number": 1}, + {"filters": 64, "kernel_size": 3, "number": 4}, + {"filters": 128, "kernel_size": 2, "number": 2}, + {"filters": 128, "kernel_size": 3, "number": 1}, + ] + kwargs["batchnorm"] = batchnorm + kwargs["attention_mechanism"] = attention["mechanism"] + if "reduction_ratio" in attention: + kwargs["attention_reduction_ratio"] = attention["reduction_ratio"] + + for task in tasks: + keras_wrapper = KerasSingleCNN( + input_shape=common_config["input_shape_keras"], + tasks=[task], + **kwargs + ) + torch_wrapper = PyTorchSingleCNN( + input_shape=common_config["input_shape_pytorch"], + tasks=[task], + **kwargs + ) + # Collect all the layers from the Keras-based model + keras_layers = [ + l.get_weights()[0].shape + for l in keras_wrapper.backbone_model.layers + if isinstance(l, (keras.layers.Conv2D, keras.layers.Dense, keras.layers.BatchNormalization)) + ] + # Collect also the dense layers from the head + keras_layers.extend([ + l.get_weights()[0].shape + for l in keras_wrapper.model.layers + if isinstance(l, keras.layers.Dense) and task in l.name + ] + ) + # Collect all the layers from the PyTorch-based model + # Note different shapes (PyTorch format: Out, In, H, W -> Keras: H, W, In, Out) + torch_layers = [] + for m in torch_wrapper.backbone_model.modules(): + if isinstance(m, torch.nn.Conv2d): + # PyTorch format: (Out, In, H, W) -> Keras format: (H, W, In, Out) + torch_layers.append( + (m.weight.shape[2], m.weight.shape[3], m.weight.shape[1], m.weight.shape[0]) + ) + elif isinstance(m, torch.nn.Linear): + torch_layers.append((m.weight.shape[1], m.weight.shape[0])) + elif isinstance(m, torch.nn.BatchNorm2d): + # BatchNorm parameters are 1D (gamma/beta/running stats) + torch_layers.append((m.weight.shape[0],)) + # Collect also the linear layers from the head + internal_key = torch_wrapper.logits_head._task_mapping[task] + p_head_module = torch_wrapper.logits_head.heads[internal_key] + torch_layers.extend([(m.weight.shape[1], m.weight.shape[0]) for m in p_head_module if isinstance(m, torch.nn.Linear)]) + # Assert structural length and individual weight shape alignment + assert len(keras_layers) == len(torch_layers), ( + f"Layer count mismatch: Keras has {len(keras_layers)}, PyTorch has {len(torch_layers)}" + ) + for idx, (k_shape, p_shape) in enumerate(zip(keras_layers, torch_layers)): + assert k_shape == p_shape, ( + f"Shape mismatch at layer {idx}: Keras shape {k_shape} vs PyTorch mapped shape {p_shape}" + ) + + +@pytest.mark.parametrize("block_type", ["basic", "bottleneck"]) +@pytest.mark.parametrize( + "first_layers", + [ + {"init_layer": None, "init_max_pool": None}, + {"init_layer": {'filters': 8, 'kernel_size': 7, 'strides': 2}, "init_max_pool": {'size': 3, 'strides': 2}}, + ], +) +@pytest.mark.parametrize( + "attention", + [ + {"mechanism": None}, + {"mechanism": "Channel-SE", "reduction_ratio": 8}, + {"mechanism": "Spatial-SE"}, + {"mechanism": "Dual-SE", "reduction_ratio": 32}, + ], +) +def test_ResNet_model_structure_parity(common_config, block_type, first_layers, attention): + """Verify that Keras and PyTorch models have matching layer counts and weight shapes.""" + tasks = common_config["tasks"] + kwargs = common_config["kwargs"].copy() + kwargs["init_layer"] = first_layers["init_layer"] + kwargs["init_max_pool"] = first_layers["init_max_pool"] + kwargs["residual_block_type"] = block_type + kwargs["architecture"] = [ + {"filters": 16, "blocks": 2}, + {"filters": 48, "blocks": 3}, + {"filters": 48, "blocks": 4}, + {"filters": 96, "blocks": 2}, + ] + kwargs["attention_mechanism"] = attention["mechanism"] + if "reduction_ratio" in attention: + kwargs["attention_reduction_ratio"] = attention["reduction_ratio"] + + for task in tasks: + keras_wrapper = KerasResNet( + input_shape=common_config["input_shape_keras"], + tasks=[task], + **kwargs + ) + torch_wrapper = PyTorchResNet( + input_shape=common_config["input_shape_pytorch"], + tasks=[task], + **kwargs + ) + # 1. Collect Keras backbone weight shapes (Conv2D / Dense) + keras_layers = [ + l.get_weights()[0].shape + for l in keras_wrapper.backbone_model.layers + if hasattr(l, "weights") and l.weights and isinstance(l, (keras.layers.Conv2D, keras.layers.Dense)) + ] + + # 2. Collect PyTorch backbone weight shapes in strict sequential block order + torch_layers = [] + # Catch initial standalone stem conv if present + if hasattr(torch_wrapper.backbone_model, "init_layer") and isinstance(torch_wrapper.backbone_model.init_layer, nn.Conv2d): + c = torch_wrapper.backbone_model.init_layer + torch_layers.append((c.kernel_size[0], c.kernel_size[1], c.in_channels, c.out_channels)) + # Iterate through stages and blocks sequentially + for module in torch_wrapper.backbone_model.modules(): + # Detect both BasicBlock and BottleneckBlock + if hasattr(module, "conv1") and hasattr(module, "conv2"): + # conv1 + c1 = module.conv1 + torch_layers.append((c1.kernel_size[0], c1.kernel_size[1], c1.in_channels, c1.out_channels)) + # conv2 + c2 = module.conv2 + torch_layers.append((c2.kernel_size[0], c2.kernel_size[1], c2.in_channels, c2.out_channels)) + # conv3 (Bottleneck only) + if hasattr(module, "conv3"): + c3 = module.conv3 + torch_layers.append((c3.kernel_size[0], c3.kernel_size[1], c3.in_channels, c3.out_channels)) + # Attention layer (if present) + if hasattr(module, "attn_layer") and module.attn_layer is not None: + # Inspect linear/conv layers inside the SE block (e.g. Channel-SE / Dual-SE / Spatial-SE) + for attn_submodule in module.attn_layer.modules(): + if isinstance(attn_submodule, nn.Linear): + torch_layers.append((attn_submodule.weight.shape[1], attn_submodule.weight.shape[0])) + elif isinstance(attn_submodule, nn.Conv2d): + torch_layers.append(( + attn_submodule.kernel_size[0], + attn_submodule.kernel_size[1], + attn_submodule.in_channels, + attn_submodule.out_channels, + )) + # shortcut projection (if present and not Identity) + if isinstance(module.shortcut, nn.Conv2d): + sc = module.shortcut + torch_layers.append((sc.kernel_size[0], sc.kernel_size[1], sc.in_channels, sc.out_channels)) + # Catch initial stem Conv2d (if present as direct child of backbone_model) + elif isinstance(module, nn.Conv2d) and module in torch_wrapper.backbone_model.children(): + torch_layers.insert(0, (module.kernel_size[0], module.kernel_size[1], module.in_channels, module.out_channels)) + + # 3. Append task heads + keras_layers.extend([ + l.get_weights()[0].shape + for l in keras_wrapper.model.layers + if isinstance(l, keras.layers.Dense) and task in l.name + ]) + + internal_key = torch_wrapper.logits_head._task_mapping[task] + p_head_module = torch_wrapper.logits_head.heads[internal_key] + torch_layers.extend([ + (m.weight.shape[1], m.weight.shape[0]) + for m in p_head_module + if isinstance(m, nn.Linear) + ]) + # Assert structural length and individual weight shape alignment + assert len(keras_layers) == len(torch_layers), ( + f"Layer count mismatch: Keras has {len(keras_layers)}, PyTorch has {len(torch_layers)}" + ) + for idx, (k_shape, p_shape) in enumerate(zip(keras_layers, torch_layers)): + assert k_shape == p_shape, ( + f"Shape mismatch at layer {idx}: Keras shape {k_shape} vs PyTorch mapped shape {p_shape}" + ) diff --git a/ctlearn/tools/__init__.py b/ctlearn/tools/__init__.py index d54df32d..8b102ffa 100644 --- a/ctlearn/tools/__init__.py +++ b/ctlearn/tools/__init__.py @@ -2,11 +2,15 @@ """ from .train_model import TrainCTLearnModel +from .keras.train_model import TrainCTLearnKerasModel +from .pytorch.train_model import TrainCTLearnPyTorchModel from .predict_LST1 import LST1PredictionTool from .predict_model import MonoPredictCTLearnModel, StereoPredictCTLearnModel __all__ = [ "TrainCTLearnModel", + "TrainCTLearnKerasModel", + "TrainCTLearnPyTorchModel", "LST1PredictionTool", "MonoPredictCTLearnModel", "StereoPredictCTLearnModel" diff --git a/ctlearn/tools/keras/__init__.py b/ctlearn/tools/keras/__init__.py new file mode 100644 index 00000000..3896413e --- /dev/null +++ b/ctlearn/tools/keras/__init__.py @@ -0,0 +1,6 @@ +"""ctlearn Keras training command line tool. +""" + +from .train_model import TrainCTLearnKerasModel + +__all__ = ["TrainCTLearnKerasModel"] \ No newline at end of file diff --git a/ctlearn/tools/keras/train_model.py b/ctlearn/tools/keras/train_model.py new file mode 100644 index 00000000..814d02da --- /dev/null +++ b/ctlearn/tools/keras/train_model.py @@ -0,0 +1,254 @@ +""" +Tool to train a Keras-based ``CTLearnModel``on R1/DL1a data using the ``DLDataReader`` and ``DLDataLoader``. +""" + +__all__ = ["TrainCTLearnKerasModel"] + +import atexit +import tensorflow as tf +import keras + +from ctlearn.core.keras.sequence import KerasSequence +from ctlearn.core.model import CTLearnModel +from ctlearn.tools.train_model import TrainCTLearnModel + + +class TrainCTLearnKerasModel(TrainCTLearnModel): + """ + Tool to train a ``~ctlearn.core.model.CTLearnModel`` Keras-based model on R1/DL1a data. + + 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 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" \\ + --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 Keras-based CTLearn model for the regression of the primary particle energy: + > ctlearn-train-keras-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 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/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/direction/ \\ + --reco cameradirection \\ + + 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/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/direction/ \\ + --reco skydirection \\ + """ + + 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) + + # Init the DLDataLoader for the + self.training_loader = KerasSequence( + DLDataReader=self.dl1dh_reader, + indices=self.training_indices, + tasks=self.reco_tasks, + 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 = KerasSequence( + DLDataReader=self.dl1dh_reader, + indices=self.validation_indices, + tasks=self.reco_tasks, + 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, + ) + + # Set up the callbacks + monitor = "val_loss" + monitor_mode = "min" + # Model checkpoint callback + model_path = f"{self.output_dir}/ctlearn_model.keras" + model_checkpoint_callback = keras.callbacks.ModelCheckpoint( + filepath=model_path, + monitor=monitor, + verbose=1, + mode=monitor_mode, + save_best_only=self.save_best_validation_only, + ) + # Tensorboard callback + tensorboard_callback = keras.callbacks.TensorBoard( + log_dir=self.output_dir, histogram_freq=1 + ) + # CSV logger callback + 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, + ] + + if self.early_stopping is not None: + # EarlyStopping callback + 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"], + ) + self.callbacks.append(early_stopping_callback) + + # Learning rate reducing callback + if self.lr_reducing is not None: + lr_reducing_callback = keras.callbacks.ReduceLROnPlateau( + monitor=monitor, + factor=self.lr_reducing["factor"], + patience=self.lr_reducing["patience"], + mode=monitor_mode, + verbose=1, + min_delta=self.lr_reducing["min_delta"], + min_lr=self.lr_reducing["min_lr"], + ) + self.callbacks.append(lr_reducing_callback) + + def start(self): + + # Open a strategy scope. + with self.strategy.scope(): + # Construct the model + self.log.info("Setting up the Keras model.") + self.model = CTLearnModel.from_name( + f"Keras{self.model_type}", + input_shape=self.training_loader.input_shape, + tasks=self.reco_tasks, + parent=self, + ).model + + # Select optimizer with appropriate arguments + optimizers = { + "Adadelta": lambda: keras.optimizers.Adadelta(learning_rate=self.learning_rate), + "Adam": lambda: keras.optimizers.Adam( + learning_rate=self.learning_rate, epsilon=self.adam_epsilon + ), + "RMSProp": lambda: keras.optimizers.RMSprop(learning_rate=self.learning_rate), + "SGD": lambda: keras.optimizers.SGD(learning_rate=self.learning_rate), + } + self.opt = optimizers[self.optimizer["name"]]() + + # Get the losses and metrics for the model + losses, metrics = self._get_losses_and_mertics(self.reco_tasks) + # Compile the model + self.log.info("Compiling CTLearn model.") + self.model.compile( + optimizer=self.opt, loss=losses, metrics=metrics + ) + + # Train and evaluate the model + self.log.info("Training and evaluating...") + self.model.fit( + self.training_loader, + validation_data=self.validation_loader, + epochs=self.n_epochs, + class_weight=self.dl1dh_reader.class_weight, + callbacks=self.callbacks, + verbose=2, + ) + self.log.info("Training and evaluating finished succesfully!") + + def _get_losses_and_mertics(self, tasks): + """ + Build the fully connected head for the CTLearn model. + + Function to build the fully connected head of the CTLearn model using the specified parameters. + + Parameters + ---------- + inputs : keras.layers.Layer + Keras layer of the model. + layers : dict + Dictionary containing the number of neurons (as value) in the fully connected head for each task (as key). + tasks : list + List of tasks to build the head for. + + Returns + ------- + logits : dict + Dictionary containing the logits for each task. + """ + losses, metrics = {}, {} + if "type" in self.reco_tasks: + losses["type"] = keras.losses.CategoricalCrossentropy( + reduction="sum_over_batch_size" + ) + metrics["type"] = [ + keras.metrics.CategoricalAccuracy(name="accuracy"), + keras.metrics.AUC(name="auc"), + ] + # 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(tasks) == 1: + losses = losses["type"] + metrics = metrics["type"] + if "energy" in self.reco_tasks: + losses["energy"] = keras.losses.MeanAbsoluteError( + reduction="sum_over_batch_size" + ) + metrics["energy"] = keras.metrics.MeanAbsoluteError(name="mae_energy") + if "cameradirection" in self.reco_tasks: + losses["cameradirection"] = keras.losses.MeanAbsoluteError( + reduction="sum_over_batch_size" + ) + 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 + + +def main(): + # Run the tool + tool = TrainCTLearnKerasModel() + tool.run() + + +if __name__ == "main": + main() \ No newline at end of file diff --git a/ctlearn/tools/predict_LST1.py b/ctlearn/tools/predict_LST1.py index 58069e56..9efb0279 100644 --- a/ctlearn/tools/predict_LST1.py +++ b/ctlearn/tools/predict_LST1.py @@ -2,9 +2,15 @@ Predict the gammaness, energy and arrival direction from lstchain DL1 data. """ +__all__ = ["LST1PredictionTool"] + +import atexit import numpy as np import tables import keras +import torch +import torch.nn as nn + from astropy import units as u from astropy.coordinates import AltAz, SkyCoord from astropy.table import Table, join, setdiff, vstack @@ -46,10 +52,13 @@ DL2_SUBARRAY_GEOMETRY_GROUP, ) from ctapipe.reco.utils import add_defaults_and_meta - from ctlearn import __version__ as ctlearn_version -from ctlearn.utils import get_lst1_subarray_description, validate_trait_dict - +from ctlearn.tools.utils import ( + FrameworkType, + setup_framework, + get_lst1_subarray_description, + validate_trait_dict, +) from dl1_data_handler.image_mapper import ImageMapper from dl1_data_handler.reader import ( get_unmapped_image, @@ -87,9 +96,9 @@ class LST1PredictionTool(Tool): --LST1PredictionTool.channels=cleaned_image \\ --LST1PredictionTool.channels=cleaned_relative_peak_time \\ --LST1PredictionTool.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/direction/ctlearn_model.cpk" \\ + --type_model="/path/to/your/type/ctlearn_model(.keras/.pth)"" \\ + --energy_model="/path/to/your/energy/ctlearn_model(.keras/.pth)"" \\ + --cameradirection_model="/path/to/your/direction/ctlearn_model(.keras/.pth)"" \\ --output output.dl2.h5 \\ --overwrite \\ """ @@ -273,7 +282,6 @@ class LST1PredictionTool(Tool): ), } - classes = classes_with_traits(ImageMapper) def setup(self): @@ -286,36 +294,36 @@ def setup(self): self.image_table_path = "/dl1/event/telescope/image/LST_LSTCam" self.parameter_table_name = "/dl1/event/telescope/parameters/LST_LSTCam" self.tel_id = 1 - + # Collect all configured model path attributes + model_paths = [ + self.load_type_model_from, + self.load_energy_model_from, + self.load_cameradirection_model_from, + ] + # Reads the model paths and set up the Keras or PyTorch framework + self.framework_type, self.num_devices, self.strategy, self.device = setup_framework(model_paths) + self.log.info("Framework selected: %s", self.framework_type.value) + if self.framework_type == FrameworkType.KERAS: + atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore + self.log.info("Using Keras MirroredStrategy with %d replica(s).", self.num_devices) + elif self.framework_type == FrameworkType.PYTORCH: + if self.device == torch.device("cpu"): + self.log.info("Using PyTorch CPU device.") + else: + self.log.info("Using PyTorch GPU device(s). Count: %d", self.num_devices) + # Load the models from the specified paths + input_shape_type, self.backbone_type, self.head_type = self._load_model( + self.load_type_model_from + ) + input_shape_energy, self.backbone_energy, self.head_energy = self._load_model( + self.load_energy_model_from + ) + input_shape_direction, self.backbone_direction, self.head_direction = self._load_model( + self.load_cameradirection_model_from + ) # Get the number of rows in the table with tables.open_file(self.input_url) as input_file: self.table_length = len(input_file.get_node(self.image_table_path)) - - # Load the models from the specified paths - if self.load_type_model_from is not None: - self.log.info("Loading the type model from %s.", self.load_type_model_from) - model_type = keras.saving.load_model(self.load_type_model_from) - input_shape = model_type.input_shape[1:] - self.backbone_type, self.head_type = self._split_model(model_type) - if self.load_energy_model_from is not None: - self.log.info( - "Loading the energy model from %s.", self.load_energy_model_from - ) - model_energy = keras.saving.load_model( - self.load_energy_model_from - ) - input_shape = model_energy.input_shape[1:] - self.backbone_energy, self.head_energy = self._split_model(model_energy) - if self.load_cameradirection_model_from is not None: - self.log.info( - "Loading the cameradirection model from %s.", self.load_cameradirection_model_from - ) - model_direction = keras.saving.load_model( - self.load_cameradirection_model_from - ) - input_shape = model_direction.input_shape[1:] - self.backbone_direction, self.head_direction = self._split_model(model_direction) - # Get the SubarrayDescription of the LST-1 telescope self.subarray = get_lst1_subarray_description(focal_length_choice=self.focal_length_choice) # Write the SubarrayDescription to the output file @@ -336,13 +344,13 @@ def setup(self): parent=self, ) # Check if the input shape of the model matches the image shape of the ImageMapper - if input_shape[0] != self.image_mapper.image_shape: - raise ToolConfigurationError( - f"The input shape of the model ('{input_shape[0]}') does not match " - f"the image shape of the ImageMapper ('{self.image_mapper.image_shape}'). " - f"Use e.g. '--BilinearMapper.interpolation_image_shape={input_shape[0]}' ." - ) - + for input_shape in [input_shape_type, input_shape_energy, input_shape_direction]: + if input_shape is not None and input_shape[0] != self.image_mapper.image_shape: + raise ToolConfigurationError( + f"The input shape of the model ('{input_shape[0]}') does not match " + f"the image shape of the ImageMapper ('{self.image_mapper.image_shape}'). " + f"Use e.g. '--BilinearMapper.interpolation_image_shape={input_shape[0]}' ." + ) # Get offset and scaling of images self.transforms = {} self.transforms["image_scale"] = 0.0 @@ -352,7 +360,6 @@ def setup(self): # Get the number of rows in the table with tables.open_file(self.input_url) as input_file: img_table_v_attrs = input_file.get_node(self.image_table_path)._v_attrs - # Check the transform value used for the file compression if "CTAFIELD_3_TRANSFORM_SCALE" in img_table_v_attrs: self.transforms["image_scale"] = img_table_v_attrs[ @@ -370,7 +377,6 @@ def setup(self): ] def start(self): - all_identifiers = read_table(self.input_url, self.parameter_table_name) all_identifiers.meta = {} if self.override_obs_id is not None: @@ -542,21 +548,20 @@ def start(self): tel_altitude.extend(dl1_table["tel_alt"].data) trigger_time.extend(dl1_table["time"].mjd) if self.load_type_model_from is not None: - classification_feature_vectors = self.backbone_type.predict_on_batch(input_data) - classification_fvs.extend(classification_feature_vectors) - predict_data = self.head_type.predict_on_batch(classification_feature_vectors) - prediction.extend(predict_data[:, 1]) + fvs, preds = self._predict_batch(self.backbone_type, self.head_type, input_data) + classification_fvs.extend(fvs) + prediction.extend(preds[:, 1]) + if self.load_energy_model_from is not None: - energy_feature_vectors = self.backbone_energy.predict_on_batch(input_data) - energy_fvs.extend(energy_feature_vectors) - predict_data = self.head_energy.predict_on_batch(energy_feature_vectors) - energy.extend(predict_data.T[0]) + fvs, preds = self._predict_batch(self.backbone_energy, self.head_energy, input_data) + energy_fvs.extend(fvs) + energy.extend(preds[:, 0]) + if self.load_cameradirection_model_from is not None: - direction_feature_vectors = self.backbone_direction.predict_on_batch(input_data) - direction_fvs.extend(direction_feature_vectors) - predict_data = self.head_direction.predict_on_batch(direction_feature_vectors) - cam_coord_offset_x.extend(predict_data.T[0]) - cam_coord_offset_y.extend(predict_data.T[1]) + fvs, preds = self._predict_batch(self.backbone_direction, self.head_direction, input_data) + direction_fvs.extend(fvs) + cam_coord_offset_x.extend(preds[:, 0]) + cam_coord_offset_y.extend(preds[:, 1]) # Create the prediction tables example_identifiers = Table( @@ -887,36 +892,119 @@ def start(self): def finish(self): self.log.info("Tool is shutting down") - def _split_model(self, model): + def _predict_batch(self, backbone, head, input_data): """ - Split the model into backbone and head. + Run batch inference through a backbone and head model using Keras or PyTorch. - This method splits the model into backbone and head. The backbone is summarized - into a single layer which can be retrieved by the layer index 1. The model input - has layer index 0 and the head is the rest of the model with layer index 2 and above. + Dispatches model execution based on ``self.framework_type``. For PyTorch, + converts input data to Tensors, transposes to channels-first format (NCHW), + places inputs on the targeted device, runs inference without gradient tracking + (``torch.no_grad()``), and returns NumPy outputs on CPU. - Parameters: - ----------- - model : keras.Model - Keras model to split into backbone and head. + Parameters + ---------- + backbone : keras.Model or torch.nn.Module + The feature extractor network that processes raw input data into + feature representations. + head : keras.Model or torch.nn.Module + The task-specific head network that computes final predictions + from feature vectors. + input_data : np.ndarray + Batch of image data following the Keras convention. - Returns: - -------- - backbone : keras.Model - Backbone model of the original model. - head : keras.Model - Head model of the original model. + Returns + ------- + fvs : np.ndarray + Extracted feature vectors for the batch. + preds : np.ndarray + Model output predictions for the batch. """ - # Get the backbone model which is the second layer of the model - backbone = model.get_layer(index=1) - # Create a new head model with the same layers as the original model. - # The output of the backbone model is the input of the head model. - backbone_output_shape = keras.Input(model.layers[2].input.shape[1:]) - x = backbone_output_shape - for layer in model.layers[2:]: - x = layer(x) - head = keras.Model(inputs=backbone_output_shape, outputs=x) - return backbone, head + if self.framework_type == FrameworkType.KERAS: + fvs = backbone.predict_on_batch(input_data) + preds = head.predict_on_batch(fvs) + return fvs, preds + if self.framework_type == FrameworkType.PYTORCH: + tensor_data = torch.from_numpy(input_data).float() + # Convert Channels-Last (N, H, W, C) to Channels-First (N, C, H, W) + if tensor_data.ndim == 4: + tensor_data = tensor_data.permute(0, 3, 1, 2) + tensor_data = tensor_data.to(self.device) + + def _to_numpy(val): + if isinstance(val, dict): + # Extract the tensor value from the dictionary (e.g., {'energy': tensor(...)}) + val = next(iter(val.values())) + if isinstance(val, torch.Tensor): + return val.cpu().numpy() + return np.asarray(val) + + with torch.no_grad(): + fvs = backbone(tensor_data) + preds = head(fvs) + return fvs.cpu().numpy(), _to_numpy(preds) + + def _load_model(self, path): + """ + Load a trained model and extract its input shape, backbone, and head. + + Supports both Keras and PyTorch frameworks based on ``self.framework_type``. + For PyTorch models, the model is set to evaluation mode (``.eval()``) and its + ``backbone`` and ``head`` submodules are extracted. + + Parameters + ---------- + path : str, pathlib.Path, or None + Path to the saved model file. If ``None``, the function immediately returns ``None`` elements. + + Returns + ------- + input_shape : tuple of int or None + Input feature tensor shape (excluding batch dimension), or ``None`` if ``path`` + is ``None`` or framework is PyTorch. + backbone : keras.Model, torch.nn.Module, or None + Feature extraction backbone of the model, or ``None`` if ``path`` is ``None``. + head : keras.Model, torch.nn.Module, or None + Output prediction head of the model, or ``None`` if ``path`` is ``None``. + + Raises + ------ + TypeError + If the loaded PyTorch object is not an instance of ``torch.nn.Module`` + (e.g., when a ``state_dict`` is loaded instead of the full model). + """ + if path is None: + return None, None, None + self.log.info("Loading the model from %s.", path) + + def _split_keras_model(model): + """Split the Keras model into backbone and head.""" + # Get the backbone model which is the second layer of the model + backbone = model.get_layer(index=1) + # Create a new head model with the same layers as the original model. + # The output of the backbone model is the input of the head model. + backbone_output_shape = keras.Input(model.layers[2].input.shape[1:]) + x = backbone_output_shape + for layer in model.layers[2:]: + x = layer(x) + head = keras.Model(inputs=backbone_output_shape, outputs=x) + return backbone, head + + if self.framework_type == FrameworkType.KERAS: + model = keras.saving.load_model(path) + input_shape = model.input_shape[1:] + backbone, head = _split_keras_model(model) + return input_shape, backbone, head + if self.framework_type == FrameworkType.PYTORCH: + model = torch.load(path, map_location=self.device, weights_only=False) + if not isinstance(model, nn.Module): + raise TypeError( + f"Expected a PyTorch 'nn.Module' object at '{path}', " + f"but got '{type(model).__name__}'. " + "Ensure the model was saved via 'torch.save(model, path)' " + "rather than 'torch.save(model.state_dict(), path)'." + ) + model.eval() + return None, model.backbone, model.head def _create_nan_table(self, nonexample_identifiers, columns, shapes): """ @@ -952,4 +1040,4 @@ def main(): if __name__ == "main": - main() + main() \ No newline at end of file diff --git a/ctlearn/tools/predict_model.py b/ctlearn/tools/predict_model.py index acfa39be..6e0a3e25 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -1,22 +1,28 @@ """ -Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``DLDataLoader``. +Tools to predict the gammaness, energy and arrival direction in monoscopic and stereoscopic mode using ``CTLearnModel`` on R1/DL1 data using the ``DLDataReader`` and ``KerasSequence``/``PyTorchDataset``. """ +__all__ = [ + "PredictCTLearnModel", + "MonoPredictCTLearnModel", + "StereoPredictCTLearnModel", +] + import atexit import uuid import warnings import numpy as np import tables -import tensorflow as tf import keras +import torch +import torch.nn as nn +from torch.utils.data import DataLoader from astropy import units as u -from astropy.coordinates.earth import EarthLocation from astropy.coordinates import AltAz, SkyCoord from astropy.table import ( Table, - hstack, vstack, join, setdiff, @@ -86,8 +92,13 @@ LST_EPOCH, ) from ctlearn import __version__ as ctlearn_version -from ctlearn.core.loader import DLDataLoader -from ctlearn.utils import validate_trait_dict +from ctlearn.core.keras.sequence import KerasSequence +from ctlearn.core.pytorch.dataset import PyTorchDataset +from ctlearn.tools.utils import ( + FrameworkType, + setup_framework, + validate_trait_dict, +) # Convienient constants for column names and table keys SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] @@ -118,11 +129,6 @@ DataLevel.DL2: DL2_SUBARRAY_GROUP, } - -class CannotPredict(OSError): - """Raised when trying to predict an incompatible file""" - - class PredictCTLearnModel(Tool): """ Base tool to predict the gammaness, energy and arrival direction from R1/DL1 data using CTLearn models. @@ -130,7 +136,8 @@ class PredictCTLearnModel(Tool): This class handles the prediction of the gammaness, energy and arrival direction from pixel-wise image or waveform data. It also supports the extraction of the feature vectors from the backbone submodel to store them in the output file. The input data is loaded from the input url using the - ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.loader.DLDataLoader``. + ``~dl1_data_handler.reader.DLDataReader`` and ``~ctlearn.core.keras.sequence.KerasSequence`` or + ``~ctlearn.core.pytorch.dataset.PyTorchDataset``. The prediction is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. The ``start`` method is implemented in the subclasses to handle the prediction for mono and stereo mode. @@ -156,14 +163,14 @@ class PredictCTLearnModel(Tool): prefix : str Name of the reconstruction algorithm used to generate the dl2 data. load_type_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the classification of the primary particle type. + Path to a Keras or PyTorch model file for the classification of the primary particle type. load_energy_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression of the primary particle energy. + Path to a Keras or PyTorch model file for the regression of the primary particle energy. load_cameradirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression + Path to a Keras or PyTorch model file for the regression of the primary particle arrival direction based on camera coordinate offsets. load_skydirection_model_from : pathlib.Path - Path to a Keras model file (Keras3) for the regression + Path to a Keras or PyTorch model file for the regression of the primary particle arrival direction based on spherical coordinate offsets. output_path : pathlib.Path Output path to save the dl2 prediction results. @@ -171,8 +178,8 @@ class PredictCTLearnModel(Tool): Verbosity mode of Keras during the prediction. strategy : tf.distribute.Strategy MirroredStrategy to distribute the prediction. - data_loader : ctlearn.core.loader.DLDataLoader - DLDataLoader object to load the data. + data_loader : ctlearn.core.keras.sequence.KerasSequence + KerasSequence object to load the data. indices : list of int List of indices for the data loaders. batch_size : int @@ -274,8 +281,8 @@ class PredictCTLearnModel(Tool): load_type_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the classification " - "of the primary particle type." + "Path to a Keras or PyTorch model file for the classification of the primary particle type. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -286,8 +293,8 @@ class PredictCTLearnModel(Tool): load_energy_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the regression " - "of the primary particle energy." + "Path to a Keras or PyTorch model file for the regression of the primary particle energy. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -298,8 +305,9 @@ class PredictCTLearnModel(Tool): load_cameradirection_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the reconstruction " - "of the primary particle arrival direction based on camera coordinate offsets." + "Path to a Keras or PyTorch model file for the reconstruction " + "of the primary particle arrival direction based on camera coordinate offsets. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -310,8 +318,9 @@ class PredictCTLearnModel(Tool): load_skydirection_model_from = Path( default_value=None, help=( - "Path to a Keras model file (Keras3) for the reconstruction " - "of the primary particle arrival direction based on spherical coordinate offsets." + "Path to a Keras or PyTorch model file for the reconstruction " + "of the primary particle arrival direction based on spherical coordinate offsets. " + "Requires Keras/PyTorch consistency with other model paths." ), allow_none=True, exists=True, @@ -331,17 +340,6 @@ class PredictCTLearnModel(Tool): help="Output path to save the dl2 prediction results", ).tag(config=True) - keras_verbose = Int( - default_value=1, - min=0, - max=2, - allow_none=False, - help=( - "Verbosity mode of Keras during the prediction: " - "0 = silent, 1 = progress bar, 2 = one line per call." - ), - ).tag(config=True) - aliases = { ("i", "input_url"): "PredictCTLearnModel.input_url", ("t", "type_model"): "PredictCTLearnModel.load_type_model_from", @@ -432,11 +430,24 @@ def setup(self): self.output_path, dl2_subarray=False, dl2_telescope=False, parent=self ) as merger: merger(self.input_url) - # 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) - + # Collect all configured model path attributes + model_paths = [ + self.load_type_model_from, + self.load_energy_model_from, + self.load_cameradirection_model_from, + self.load_skydirection_model_from, + ] + # Reads the model paths and set up the Keras or PyTorch framework + self.framework_type, self.num_devices, self.strategy, self.device = setup_framework(model_paths) + self.log.info("Framework selected: %s", self.framework_type.value) + if self.framework_type == FrameworkType.KERAS: + atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore + self.log.info("Using Keras MirroredStrategy with %d replica(s).", self.num_devices) + elif self.framework_type == FrameworkType.PYTORCH: + if self.device == torch.device("cpu"): + self.log.info("Using PyTorch CPU device.") + else: + self.log.info("Using PyTorch GPU device(s). Count: %d", self.num_devices) # Set up the data reader self.log.info("Loading data reader:") self.log.info("For a large dataset, this may take a while...") @@ -455,7 +466,7 @@ def setup(self): # Set the indices for the data loaders self.indices = list(range(self.dl1dh_reader._get_n_events())) self.last_batch_size = len(self.indices) % ( - self.batch_size * self.strategy.num_replicas_in_sync + self.batch_size * self.num_devices ) # Ensure subarray consistency in the output file self._ensure_subarray_consistency() @@ -676,8 +687,15 @@ def deduplicate_first_valid( return unique(t, keys=list(keys), keep="first") def _predict_with_model(self, model_path): + """ Select the framework to load and predict the data. """ + if self.framework_type == FrameworkType.KERAS: + return self._predict_with_keras_model(model_path) + elif self.framework_type == FrameworkType.PYTORCH: + return self._predict_with_pytorch_model(model_path) + + def _predict_with_keras_model(self, model_path): """ - Load and predict with a CTLearn model. + Load and predict with a CTLearn Keras-based model. Load a model from the specified path and predict the data using the loaded model. If a last batch loader is provided, predict the last batch and stack the results. @@ -694,13 +712,13 @@ def _predict_with_model(self, model_path): feature_vectors : np.ndarray Feature vectors extracted from the backbone model. """ - # Create a new DLDataLoader for each task - # It turned out to be more robust to initialize the DLDataLoader separately. - data_loader = DLDataLoader( + # Create a new KerasSequence for each task + # It turned out to be more robust to initialize the KerasSequence separately. + data_loader = KerasSequence( self.dl1dh_reader, self.indices, tasks=[], - batch_size=self.batch_size * self.strategy.num_replicas_in_sync, + batch_size=self.batch_size * self.num_devices, sort_by_intensity=self.sort_by_intensity, stack_telescope_images=self.stack_telescope_images, ) @@ -711,7 +729,7 @@ def _predict_with_model(self, model_path): data_loader_last_batch = None if self.last_batch_size > 0: last_batch_indices = self.indices[-self.last_batch_size :] - data_loader_last_batch = DLDataLoader( + data_loader_last_batch = KerasSequence( self.dl1dh_reader, last_batch_indices, tasks=[], @@ -740,7 +758,7 @@ def _predict_with_model(self, model_path): # Apply the backbone model with the data loader to retrieve the feature vectors try: feature_vectors = backbone_model.predict( - data_loader, verbose=self.keras_verbose + data_loader ) except ValueError as err: if str(err).startswith("Input 0 of layer"): @@ -754,14 +772,14 @@ def _predict_with_model(self, model_path): predict_data = Table( { prediction_colname: head.predict( - feature_vectors, verbose=self.keras_verbose + feature_vectors ) } ) # Predict the last batch and stack the results to the prediction data if data_loader_last_batch is not None: feature_vectors_last_batch = backbone_model.predict( - data_loader_last_batch, verbose=self.keras_verbose + data_loader_last_batch ) feature_vectors = np.concatenate( (feature_vectors, feature_vectors_last_batch) @@ -772,8 +790,7 @@ def _predict_with_model(self, model_path): Table( { prediction_colname: head.predict( - feature_vectors_last_batch, - verbose=self.keras_verbose, + feature_vectors_last_batch ) } ), @@ -782,7 +799,7 @@ def _predict_with_model(self, model_path): else: # Predict the data using the loaded model try: - predict_data = model.predict(data_loader, verbose=self.keras_verbose) + predict_data = model.predict(data_loader) except ValueError as err: if str(err).startswith("Input 0 of layer"): raise ToolConfigurationError( @@ -803,7 +820,7 @@ def _predict_with_model(self, model_path): # Predict the last batch and stack the results to the prediction data if data_loader_last_batch is not None: predict_data_last_batch = model.predict( - data_loader_last_batch, verbose=self.keras_verbose + data_loader_last_batch ) if model.layers[-1].name == "type": predict_data_last_batch = Table( @@ -814,6 +831,80 @@ def _predict_with_model(self, model_path): predict_data = vstack([predict_data, predict_data_last_batch]) return predict_data, feature_vectors + + def _predict_with_pytorch_model(self, model_path): + """ + Load and predict with a CTLearn PyTorch-based model. + + Parameters + ---------- + model_path : str + Path to a PyTorch model checkpoint/file. + + Returns + ------- + predict_data : astropy.table.Table + Table containing the prediction results. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + # Load the model and make sure it is not a state dict + model = torch.load(model_path, map_location=self.device, weights_only=False) + if not isinstance(model, nn.Module): + raise TypeError( + f"Expected a PyTorch 'nn.Module' object at '{model_path}', " + f"but got '{type(model).__name__}'. " + "Ensure the model was saved via 'torch.save(model, path)' rather than 'torch.save(model.state_dict(), path)'." + ) + model.eval() + # PyTorch DataLoaders natively handle drop_last=False, + # so we don't need a separate generator for incomplete batches. + # Replace 'PyTorchDataset' with your actual PyTorch Dataset implementation + dataset = PyTorchDataset( + self.dl1dh_reader, + self.indices, + tasks=[], + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + data_loader = DataLoader( + dataset, + batch_size=self.batch_size * self.num_devices, + shuffle=False, + pin_memory=torch.cuda.is_available(), + drop_last=False + ) + predictions_dict, feature_vectors_list = {}, [] + with torch.no_grad(): + for batch in data_loader: # Iterate over DataLoader without subscripting + inputs = batch[0].to(self.device) + predictions, features = model(inputs) + if self.dl1_features: + feature_vectors_list.append(features.cpu().numpy()) + # Standardize model outputs into predictions_dict using self.heads_dict + if isinstance(predictions, torch.Tensor): + task_name = list(model.head.heads_dict.keys())[0] + predictions_dict.setdefault(task_name, []).append(predictions.cpu().numpy()) + elif isinstance(predictions, dict): + for task_name, output_tensor in predictions.items(): + predictions_dict.setdefault(task_name, []).append( + output_tensor.cpu().numpy() + ) + # Concatenate batched prediction arrays and format into Astropy Table + predict_data = Table( + { + task_name: np.concatenate(batches, axis=0) + for task_name, batches in predictions_dict.items() + } + ) + # Concatenate features if extracted + feature_vectors = ( + np.concatenate(feature_vectors_list, axis=0) + if self.dl1_features and feature_vectors_list + else None + ) + return predict_data, feature_vectors + def _predict_particletype(self, example_identifiers): """ Predict the classification of the primary particle type. @@ -1332,9 +1423,9 @@ class MonoPredictCTLearnModel(PredictCTLearnModel): --DLImageReader.channels=cleaned_image \\ --DLImageReader.channels=cleaned_relative_peak_time \\ --DLImageReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model.cpk" \\ + --type_model="/path/to/your/mono/type/ctlearn_model(.keras/.pth)" \\ + --energy_model="/path/to/your/mono/energy/ctlearn_model(.keras/.pth)" \\ + --cameradirection_model="/path/to/your/mono/cameradirection/ctlearn_model(.keras/.pth)" \\ --dl1-features \\ --no-dl1-images \\ --no-true-images \\ @@ -1346,9 +1437,9 @@ class MonoPredictCTLearnModel(PredictCTLearnModel): --PredictCTLearnModel.dl1dh_reader_type=DLWaveformReader \\ --DLWaveformReader.sequnce_length=20 \\ --DLWaveformReader.image_mapper_type=BilinearMapper \\ - --type_model="/path/to/your/mono_waveform/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model.cpk" \\ - --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model.cpk" \\ + --type_model="/path/to/your/mono_waveform/type/ctlearn_model(.keras/.pth)" \\ + --energy_model="/path/to/your/mono_waveform/energy/ctlearn_model(.keras/.pth)" \\ + --cameradirection_model="/path/to/your/mono_waveform/cameradirection/ctlearn_model(.keras/.pth)" \\ --no-r0-waveforms \\ --no-r1-waveforms \\ --no-dl1-images \\ @@ -1992,9 +2083,9 @@ class StereoPredictCTLearnModel(PredictCTLearnModel): --DLImageReader.mode=stereo \\ --DLImageReader.min_telescopes=2 \\ --PredictCTLearnModel.stack_telescope_images=True \\ - --type_model="/path/to/your/stereo/type/ctlearn_model.cpk" \\ - --energy_model="/path/to/your/stereo/energy/ctlearn_model.cpk" \\ - --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model.cpk" \\ + --type_model="/path/to/your/stereo/type/ctlearn_model(.keras/.pth)" \\ + --energy_model="/path/to/your/stereo/energy/ctlearn_model(.keras/.pth)" \\ + --skydirection_model="/path/to/your/stereo/skydirection/ctlearn_model(.keras/.pth)" \\ --output output.dl2.h5 \\ """ @@ -2366,4 +2457,4 @@ def stereo_tool(): mono_tool() if __name__ == "stereo_tool": - stereo_tool() + stereo_tool() \ No newline at end of file diff --git a/ctlearn/tools/pytorch/__init__.py b/ctlearn/tools/pytorch/__init__.py new file mode 100644 index 00000000..6bf284bd --- /dev/null +++ b/ctlearn/tools/pytorch/__init__.py @@ -0,0 +1,6 @@ +"""ctlearn PyTorch training command line tool. +""" + +from .train_model import TrainCTLearnPyTorchModel + +__all__ = ["TrainCTLearnPyTorchModel"] \ No newline at end of file diff --git a/ctlearn/tools/pytorch/train_model.py b/ctlearn/tools/pytorch/train_model.py new file mode 100644 index 00000000..d80affb9 --- /dev/null +++ b/ctlearn/tools/pytorch/train_model.py @@ -0,0 +1,484 @@ +""" +Tool to train a PyTorch-based ``CTLearnModel`` on R1/DL1a data using the ``DLDataReader`` and ``PyTorchDataLoader``. +""" + +__all__ = ["TrainCTLearnPyTorchModel"] + +import os +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from torch.utils.tensorboard import SummaryWriter +import torchmetrics + +from ctlearn.core.pytorch.dataset import PyTorchDataset +from ctlearn.core.model import CTLearnModel +from ctlearn.tools.train_model import TrainCTLearnModel + + +class TrainCTLearnPyTorchModel(TrainCTLearnModel): + """ + Tool to train a ``~ctlearn.core.model.CTLearnModel`` PyTorch-based model on R1/DL1a data. + + The tool trains a CTLearn PyTorch-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-pytorch-model" + description = __doc__ + + examples = """ + To train a PyTorch-based CTLearn 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 PyTorch-based CTLearn 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 PyTorch-based CTLearn 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 PyTorch-based CTLearn 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 \\ + """ + + def setup_framework(self): + # Determine available hardware device (Multi-GPU / Single GPU / CPU) + if torch.cuda.is_available(): + self.device = torch.device("cuda") + self.num_devices = torch.cuda.device_count() + self.log.info("Using CUDA device(s). Available GPUs: %d", self.num_devices) + else: + self.device = torch.device("cpu") + self.num_devices = 1 + self.log.info("Using CPU device.") + + + # Set seed globally for reproducibility across random operations + torch.manual_seed(self.random_seed) + # Create a dedicated Generator for the DataLoader + g = torch.Generator() + g.manual_seed(self.random_seed) + + # Init the PyTorchDataLoader + self.training_dataset = PyTorchDataset( + DLDataReader=self.dl1dh_reader, + indices=self.training_indices, + tasks=self.reco_tasks, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + self.training_loader = DataLoader( + dataset=self.training_dataset, + batch_size=self.batch_size * self.num_devices, + shuffle=True, # Enables shuffling + generator=g, # Controls the shuffling seed deterministically + pin_memory=torch.cuda.is_available() # Accelerates memory copy from host CPU to GPU + ) + self.validation_dataset = PyTorchDataset( + DLDataReader=self.dl1dh_reader, + indices=self.validation_indices, + tasks=self.reco_tasks, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + self.validation_loader = DataLoader( + dataset=self.validation_dataset, + batch_size=self.batch_size * self.num_devices, + shuffle=False, # Disables shuffling + pin_memory=torch.cuda.is_available() # Accelerates memory copy from host CPU to GPU + ) + + # Set up TensorBoard writers for train and validation and CSV logging path + self.train_writer = SummaryWriter(log_dir=os.path.join(self.output_dir, "train")) + self.val_writer = SummaryWriter(log_dir=os.path.join(self.output_dir, "validation")) + self.csv_log_path = os.path.join(self.output_dir, "training_log.csv") + + # Initialize TorchMetrics according to reconstruction tasks + self.train_metrics = self._get_task_metrics() + self.val_metrics = self._get_task_metrics() + + # Build CSV dynamic header matching Keras logger format + self.csv_headers = ["epoch"] + if "type" in self.reco_tasks: + self.csv_headers.extend(["accuracy", "auc"]) + if "energy" in self.reco_tasks: + self.csv_headers.append("mae_energy") + if "cameradirection" in self.reco_tasks: + self.csv_headers.append("mae_cameradirection") + if "skydirection" in self.reco_tasks: + self.csv_headers.append("mae_skydirection") + self.csv_headers.append("loss") + + # Add validation metrics to header + if "type" in self.reco_tasks: + self.csv_headers.extend(["val_accuracy", "val_auc"]) + if "energy" in self.reco_tasks: + self.csv_headers.append("val_mae_energy") + if "cameradirection" in self.reco_tasks: + self.csv_headers.append("val_mae_cameradirection") + if "skydirection" in self.reco_tasks: + self.csv_headers.append("val_mae_skydirection") + self.csv_headers.append("val_loss") + + # Write CSV header if file doesn't exist + if not os.path.exists(self.csv_log_path): + with open(self.csv_log_path, "w") as f: + f.write(",".join(self.csv_headers) + "\n") + + def start(self): + self.log.info("Setting up the PyTorch model.") + base_model = CTLearnModel.from_name( + f"PyTorch{self.model_type}", + input_shape=self.training_dataset.input_shape, + tasks=self.reco_tasks, + parent=self, + ).model + + base_model.to(self.device) + + if self.device.type == "cuda" and self.num_devices > 1: + self.model = nn.DataParallel(base_model) + else: + self.model = base_model + + optimizers = { + "Adadelta": lambda params: torch.optim.Adadelta(params, lr=self.learning_rate), + "Adam": lambda params: torch.optim.Adam(params, lr=self.learning_rate, eps=self.adam_epsilon), + "RMSProp": lambda params: torch.optim.RMSprop(params, lr=self.learning_rate), + "SGD": lambda params: torch.optim.SGD(params, lr=self.learning_rate), + } + self.opt = optimizers[self.optimizer["name"]](self.model.parameters()) + + # Setup Learning Rate Scheduler + self.scheduler = None + if self.lr_reducing is not None: + self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + self.opt, + mode="min", + factor=self.lr_reducing["factor"], + patience=self.lr_reducing["patience"], + threshold=self.lr_reducing["min_delta"], + min_lr=self.lr_reducing["min_lr"], + ) + + self.loss_fns = self._get_loss_functions() + + self.log.info("Training and evaluating...") + + best_val_loss = float("inf") + patience_counter = 0 + best_model_weights = None + model_save_path = os.path.join(self.output_dir, "ctlearn_model.pth") + state_dict_save_path = os.path.join(self.output_dir, "ctlearn_state_dict.pth") + + for epoch_idx in range(self.n_epochs): + train_loss, train_metric_vals = self._train_epoch() + val_loss, val_metric_vals = self._validate_epoch() + + # Record CSV Row (0-indexed epochs matching Keras CSVLogger) + row_dict = {"epoch": epoch_idx, "loss": train_loss, "val_loss": val_loss} + for k, v in train_metric_vals.items(): + row_dict[k] = v + for k, v in val_metric_vals.items(): + row_dict[f"val_{k}"] = v + + row_str = ",".join(str(row_dict[col]) for col in self.csv_headers) + with open(self.csv_log_path, "a") as f: + f.write(row_str + "\n") + + # TensorBoard metrics training and validation + self.train_writer.add_scalar("loss", train_loss, epoch_idx) + for k, v in train_metric_vals.items(): + self.train_writer.add_scalar(k, v, epoch_idx) + self.train_writer.flush() + self.val_writer.add_scalar("loss", val_loss, epoch_idx) + for k, v in val_metric_vals.items(): + self.val_writer.add_scalar(k, v, epoch_idx) + self.val_writer.flush() + + # Unwrap model if wrapped with DataParallel / DistributedDataParallel + unwrapped_model = ( + self.model.module + if isinstance(self.model, (nn.DataParallel, nn.parallel.DistributedDataParallel)) + else self.model + ) + + # Checkpoint saving & Early stopping logic + if val_loss < best_val_loss: + best_val_loss = val_loss + patience_counter = 0 + state_dict = ( + self.model.module.state_dict() + if isinstance(self.model, nn.DataParallel) + else self.model.state_dict() + ) + best_model_weights = state_dict + if self.save_best_validation_only: + # Save state_dict (tensors only) + torch.save(state_dict, state_dict_save_path) + # Save full model (nn.Module object) + torch.save(unwrapped_model, model_save_path) + else: + patience_counter += 1 + + if not self.save_best_validation_only: + state_dict = ( + self.model.module.state_dict() + if isinstance(self.model, nn.DataParallel) + else self.model.state_dict() + ) + # Save state_dict (tensors only) + torch.save(state_dict, state_dict_save_path) + # Save full model (nn.Module object) + torch.save(unwrapped_model, model_save_path) + + if self.scheduler is not None: + self.scheduler.step(val_loss) + + if self.early_stopping is not None: + if patience_counter >= self.early_stopping["patience"]: + self.log.info("Early stopping triggered at epoch %d.", epoch_idx) + if ( + self.early_stopping["restore_best_weights"] + and best_model_weights is not None + ): + unwrapped_model = ( + self.model.module + if isinstance(self.model, nn.DataParallel) + else self.model + ) + unwrapped_model.load_state_dict(best_model_weights) + break + + # Close the TensorBoard writers + self.train_writer.close() + self.val_writer.close() + self.log.info("Training and evaluating finished successfully!") + + def _train_epoch(self): + self.model.train() + for m in self.train_metrics.values(): + m.reset() + + total_loss = 0.0 + total_samples = 0 + + for batch_x, batch_y in self.training_loader: + batch_x = self._to_device(batch_x) + batch_y = self._to_device(batch_y) + + batch_size = batch_x.size(0) if isinstance(batch_x, torch.Tensor) else next(iter(batch_x.values())).size(0) + + self.opt.zero_grad() + outputs, _ = self.model(batch_x) + + loss = self._compute_combined_loss(outputs, batch_y) + loss.backward() + self.opt.step() + + total_loss += loss.item() * batch_size + total_samples += batch_size + self._update_metrics(self.train_metrics, outputs, batch_y) + + avg_loss = total_loss / total_samples + metric_results = {k: m.compute().item() for k, m in self.train_metrics.items()} + return avg_loss, metric_results + + def _validate_epoch(self): + self.model.eval() + self._reset_metrics(self.val_metrics) + + total_loss = 0.0 + total_samples = 0 + + with torch.no_grad(): + for batch_x, batch_y in self.validation_loader: + batch_x = self._to_device(batch_x) + batch_y = self._to_device(batch_y) + + # Dynamically determine batch size (works for Tensors or dicts of Tensors) + batch_size = ( + batch_x.size(0) + if isinstance(batch_x, torch.Tensor) + else next(iter(batch_x.values())).size(0) + ) + + outputs, _ = self.model(batch_x) + loss = self._compute_combined_loss(outputs, batch_y) + + total_loss += loss.item() * batch_size + total_samples += batch_size + + self._update_metrics(self.val_metrics, outputs, batch_y) + + avg_loss = total_loss / total_samples if total_samples > 0 else 0.0 + metric_results = self._compute_metrics(self.val_metrics) + return avg_loss, metric_results + + + # Helper methods to safely handle flat or task-nested metric dictionaries + def _reset_metrics(self, metrics): + for m in metrics.values(): + if isinstance(m, dict): + for sub_m in m.values(): + sub_m.reset() + else: + m.reset() + + def _compute_metrics(self, metrics): + results = {} + for k, v in metrics.items(): + if isinstance(v, dict): + for sub_k, sub_v in v.items(): + results[f"{k}_{sub_k}"] = sub_v.compute().item() + else: + results[k] = v.compute().item() + return results + + def _get_task_metrics(self): + """Instantiates TorchMetrics matching Keras metric definitions.""" + metrics = {} + if "type" in self.reco_tasks: + num_classes = getattr(self.dl1dh_reader, "num_classes", 2) + metrics["accuracy"] = torchmetrics.Accuracy( + task="multiclass" if num_classes > 2 else "binary", + num_classes=num_classes, + ).to(self.device) + metrics["auc"] = torchmetrics.AUROC( + task="multiclass" if num_classes > 2 else "binary", + num_classes=num_classes, + ).to(self.device) + if "energy" in self.reco_tasks: + metrics["mae_energy"] = torchmetrics.MeanAbsoluteError().to(self.device) + if "cameradirection" in self.reco_tasks: + metrics["mae_cameradirection"] = torchmetrics.MeanAbsoluteError().to(self.device) + if "skydirection" in self.reco_tasks: + metrics["mae_skydirection"] = torchmetrics.MeanAbsoluteError().to(self.device) + return metrics + + def _update_metrics(self, metrics, outputs, targets): + for task in self.reco_tasks: + out = outputs[task] if isinstance(outputs, dict) else outputs + tgt = targets[task] if isinstance(targets, dict) else targets + task_metrics = metrics[task] if task in metrics else metrics + + if task == "type": + # Ensure target is 1D LongTensor for classification + if tgt.ndim > 1: + tgt = tgt.squeeze(-1) + tgt = tgt.long() + + # Convert 2-class logits to class 1 probabilities for binary metrics + if out.ndim > 1 and out.shape[-1] > 1: + probs = torch.softmax(out, dim=-1) + out = probs[:, 1] + elif out.ndim > 1 and out.shape[-1] == 1: + out = torch.sigmoid(out.squeeze(-1)) + else: + if out.ndim > 1 and out.shape[-1] == 1: + out = out.squeeze(-1) + if tgt.ndim > 1 and tgt.shape[-1] == 1: + tgt = tgt.squeeze(-1) + + if isinstance(task_metrics, dict): + for metric_obj in task_metrics.values(): + metric_obj.update(out, tgt) + else: + task_metrics.update(out, tgt) + + def _compute_combined_loss(self, outputs, targets): + total_loss = 0.0 + + for task in self.reco_tasks: + task_output = outputs[task] if isinstance(outputs, dict) else outputs + task_target = targets[task] if isinstance(targets, dict) else targets + + if task == "type": + if task_target.ndim > 1: + task_target = task_target.squeeze(-1) + task_target = task_target.long() + else: + if task_output.ndim > 1 and task_output.shape[-1] == 1: + task_output = task_output.squeeze(-1) + if task_target.ndim > 1 and task_target.shape[-1] == 1: + task_target = task_target.squeeze(-1) + + task_loss = self.loss_fns[task](task_output, task_target) + total_loss += task_loss + + return total_loss + + def _get_loss_functions(self): + loss_fns = {} + if "type" in self.reco_tasks: + weight = None + if self.dl1dh_reader.class_weight is not None: + class_weights = self.dl1dh_reader.class_weight + + # Convert dict to a list ordered by class index + if isinstance(class_weights, dict): + class_weights = [class_weights[k] for k in sorted(class_weights.keys())] + + weight = torch.tensor( + class_weights, + dtype=torch.float32, + device=self.device, + ) + + loss_fns["type"] = torch.nn.CrossEntropyLoss(weight=weight) + if "energy" in self.reco_tasks: + loss_fns["energy"] = nn.L1Loss() + if "cameradirection" in self.reco_tasks: + loss_fns["cameradirection"] = nn.L1Loss() + if "skydirection" in self.reco_tasks: + loss_fns["skydirection"] = nn.L1Loss() + return loss_fns + + def _to_device(self, data): + if isinstance(data, torch.Tensor): + return data.to(self.device) + elif isinstance(data, dict): + return {k: self._to_device(v) for k, v in data.items()} + elif isinstance(data, list): + return [self._to_device(v) for v in data] + return data + + +def main(): + tool = TrainCTLearnPyTorchModel() + tool.run() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ctlearn/tools/tests/test_predict_LST1.py b/ctlearn/tools/tests/test_predict_LST1.py index 62969979..1fc02838 100644 --- a/ctlearn/tools/tests/test_predict_LST1.py +++ b/ctlearn/tools/tests/test_predict_LST1.py @@ -4,6 +4,7 @@ from ctapipe.core import run_tool from ctapipe.io import TableLoader +from ctlearn.conftest import MODEL_FILE_FORMATS from ctlearn.tools import LST1PredictionTool # Columns that should be present in the output DL2 file @@ -31,36 +32,31 @@ @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) def test_predict_mono_model_with_lst1_mock_data( - tmp_path, ctlearn_trained_dl1_mono_models, mock_lst1_dl1_file + tmp_path, ctlearn_trained_dl1_mono_models, mock_lst1_dl1_file, framework ): """ Test LST1PredictionTool using trained mono models and mock LST-1 DL1 files. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) - # Hardcopy the trained models to the model directory telescope_type = "LST" for reco_task in ["type", "energy", "cameradirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_dl1_mono_models[key], - model_dir / f"ctlearn_mono_model_{key}.keras", + model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_mono_model_{key}.keras" + model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained mono model file not found for {key}" - # Check that the mock LST1 DL1 file was created assert mock_lst1_dl1_file.exists(), "Mock LST1 DL1 file not found" - - output_file = dl2_dir / "mock_lst1_predictions.dl2.h5" - + output_file = dl2_dir / f"mock_lst1_{framework}_predictions.dl2.h5" # Build command-line arguments for LST1PredictionTool argv = [ f"--input_url={mock_lst1_dl1_file}", @@ -69,19 +65,16 @@ def test_predict_mono_model_with_lst1_mock_data( "--LST1PredictionTool.channels=cleaned_image", "--LST1PredictionTool.channels=cleaned_relative_peak_time", "--LST1PredictionTool.image_mapper_type=BilinearMapper", - f"--type_model={model_dir}/ctlearn_mono_model_{telescope_type}_type.keras", - f"--energy_model={model_dir}/ctlearn_mono_model_{telescope_type}_energy.keras", - f"--cameradirection_model={model_dir}/ctlearn_mono_model_{telescope_type}_cameradirection.keras", + f"--type_model={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--energy_model={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--cameradirection_model={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", "--dl2-telescope", "--overwrite", ] - # Run LST1PredictionTool assert run_tool(LST1PredictionTool(), argv=argv, cwd=tmp_path) == 0 - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" - # Check that the created DL2 file can be read with the TableLoader allowed_tels = [1] with TableLoader( @@ -107,7 +100,6 @@ def test_predict_mono_model_with_lst1_mock_data( assert ( tel_events[tel_id][col][0] is not np.nan ), f"{col} has NaN values in DL2 file {output_file.name}" - # Check subarray-wise data subarray_events = loader.read_subarray_events(start=0, stop=2, dl2=True) assert len(subarray_events) > 0 diff --git a/ctlearn/tools/tests/test_predict_model.py b/ctlearn/tools/tests/test_predict_model.py index 5f17fc34..1434192b 100644 --- a/ctlearn/tools/tests/test_predict_model.py +++ b/ctlearn/tools/tests/test_predict_model.py @@ -4,6 +4,7 @@ from ctapipe.core import run_tool from ctapipe.io import TableLoader +from ctlearn.conftest import MODEL_FILE_FORMATS from ctlearn.tools import MonoPredictCTLearnModel, StereoPredictCTLearnModel # Columns that should be present in the output DL2 file @@ -36,31 +37,29 @@ @pytest.mark.verifies_usecase("DPPS-UC-130-1.2") +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) def test_predict_mono_model_with_r1_waveforms( - tmp_path, ctlearn_trained_r1_mono_models, r1_gamma_file + tmp_path, ctlearn_trained_r1_mono_models, r1_gamma_file, framework ): """ Test training CTLearn mono model using the R1 gamma and proton files for all reconstruction tasks and predicting DL2 from R1 waveforms. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) # Define telescope types and their available telescopes telescope_type = "LST" available_tels = [1, 2] - # Hardcopy the trained models to the model directory for reco_task in ["type", "energy", "cameradirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_r1_mono_models[key], - model_dir / f"ctlearn_mono_model_{key}.keras", + model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_mono_model_{key}.keras" + model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained mono model file not found for {key}" # Build command-line arguments argv = [ @@ -72,7 +71,7 @@ def test_predict_mono_model_with_r1_waveforms( "--no-r1-waveforms", "--dl2-telescope", ] - output_file = dl2_dir / f"gamma_{telescope_type}_mono_from_waveforms.dl2.h5" + output_file = dl2_dir / f"gamma_{framework}_{telescope_type}_mono_from_waveforms.dl2.h5" # Run Prediction tool assert ( run_tool( @@ -80,15 +79,14 @@ def test_predict_mono_model_with_r1_waveforms( argv=argv + [ f"--output={output_file}", - f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_type.keras", - f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_energy.keras", - f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_cameradirection.keras", + f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) == 0 ) - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" # Check that the created DL2 file can be read with the TableLoader @@ -129,18 +127,17 @@ def test_predict_mono_model_with_r1_waveforms( @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) @pytest.mark.parametrize("dl2_tel_flag", ["dl2-telescope", "no-dl2-telescope"]) def test_predict_mono_model_with_dl1_images( - tmp_path, ctlearn_trained_dl1_mono_models, dl1_gamma_file, dl2_tel_flag + tmp_path, ctlearn_trained_dl1_mono_models, dl1_gamma_file, framework, dl2_tel_flag ): """ Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks and predicting DL2 from DL1 images. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) # Define telescope types and their allowed telescopes @@ -157,12 +154,12 @@ def test_predict_mono_model_with_dl1_images( # Hardcopy the trained models to the model directory for telescope_type in telescope_types.keys(): for reco_task in ["type", "energy", "cameradirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_dl1_mono_models[key], - model_dir / f"ctlearn_mono_model_{key}.keras", + model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_mono_model_{key}.keras" + model_file = model_dir / f"ctlearn_mono_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained mono model file not found for {key}" # Build command-line arguments argv = [ @@ -175,7 +172,7 @@ def test_predict_mono_model_with_dl1_images( ] for telescope_type, allowed_tels in telescope_types.items(): output_file = ( - dl2_dir / f"gamma_{dl2_tel_flag}_{telescope_type}_mono_from_images.dl2.h5" + dl2_dir / f"gamma_{dl2_tel_flag}_{telescope_type}_{framework}_mono_from_images.dl2.h5" ) # Run Prediction tool assert ( @@ -186,15 +183,14 @@ def test_predict_mono_model_with_dl1_images( f"--output={output_file}", f"--DLImageReader.allowed_tels={allowed_tels}", f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", - f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_type.keras", - f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_energy.keras", - f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{telescope_type}_cameradirection.keras", + f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_cameradirection_model_from={model_dir}/ctlearn_mono_model_{framework}_{telescope_type}_cameradirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) == 0 ) - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" # Check that the created DL2 file can be read with the TableLoader @@ -237,31 +233,29 @@ def test_predict_mono_model_with_dl1_images( @pytest.mark.verifies_usecase("DPPS-UC-130-1.2.2") +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) def test_predict_stereo_model_with_dl1_images( - tmp_path, ctlearn_trained_dl1_stereo_models, dl1_gamma_file + tmp_path, ctlearn_trained_dl1_stereo_models, dl1_gamma_file, framework ): """ Test training CTLearn stereo model using the DL1 gamma and proton files for all reconstruction tasks and predicting DL2 from DL1 images. Each test run gets its own isolated temp directories. """ - model_dir = tmp_path / "trained_models" model_dir.mkdir(parents=True, exist_ok=True) - dl2_dir = tmp_path / "dl2_output" dl2_dir.mkdir(parents=True, exist_ok=True) # Define telescope types and their available telescopes telescope_type = "MST" allowed_tels = [7, 13, 15] - # Hardcopy the trained models to the model directory for reco_task in ["type", "energy", "skydirection"]: - key = f"{telescope_type}_{reco_task}" + key = f"{framework}_{telescope_type}_{reco_task}" shutil.copy( ctlearn_trained_dl1_stereo_models[key], - model_dir / f"ctlearn_stereo_model_{key}.keras", + model_dir / f"ctlearn_stereo_model_{key}.{MODEL_FILE_FORMATS[framework]}", ) - model_file = model_dir / f"ctlearn_stereo_model_{key}.keras" + model_file = model_dir / f"ctlearn_stereo_model_{key}.{MODEL_FILE_FORMATS[framework]}" assert model_file.exists(), f"Trained stereo model file not found for {key}" # Build command-line arguments argv = [ @@ -274,7 +268,7 @@ def test_predict_stereo_model_with_dl1_images( "--no-dl1-images", "--no-true-images", ] - output_file = dl2_dir / f"gamma_{telescope_type}_stereo_from_images.dl2.h5" + output_file = dl2_dir / f"gamma_{framework}_{telescope_type}_stereo_from_images.dl2.h5" # Run Prediction tool assert ( run_tool( @@ -282,15 +276,14 @@ def test_predict_stereo_model_with_dl1_images( argv=argv + [ f"--output={output_file}", - f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_stereo_model_{telescope_type}_type.keras", - f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_stereo_model_{telescope_type}_energy.keras", - f"--PredictCTLearnModel.load_skydirection_model_from={model_dir}/ctlearn_stereo_model_{telescope_type}_skydirection.keras", + f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_type.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_energy.{MODEL_FILE_FORMATS[framework]}", + f"--PredictCTLearnModel.load_skydirection_model_from={model_dir}/ctlearn_stereo_model_{framework}_{telescope_type}_skydirection.{MODEL_FILE_FORMATS[framework]}", ], cwd=tmp_path, ) == 0 ) - # Check that the output DL2 file was created assert output_file.exists(), "Output DL2 file not created" # Check that the created DL2 file can be read with the TableLoader diff --git a/ctlearn/tools/tests/test_train_model.py b/ctlearn/tools/tests/test_train_model.py index d2291652..3466f2cf 100644 --- a/ctlearn/tools/tests/test_train_model.py +++ b/ctlearn/tools/tests/test_train_model.py @@ -3,31 +3,41 @@ import shutil from ctapipe.core import run_tool -from ctlearn.tools import TrainCTLearnModel - +from ctlearn.conftest import TRAINING_TOOLS, MODEL_FILE_FORMATS +@pytest.mark.parametrize("framework", ["Keras", "PyTorch"]) +@pytest.mark.parametrize("model", ["SingleCNN", "ResNet", "LoadedModel"]) @pytest.mark.parametrize("reco_task", ["type", "energy", "cameradirection"]) -def test_train_ctlearn_model(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) - background_dir = tmp_path / "proton_dl1" background_dir.mkdir(parents=True, exist_ok=True) - # Hardcopy DL1 gamma file to the signal directory shutil.copy(dl1_gamma_file, signal_dir) # 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 = [ f"--signal={signal_dir}", @@ -39,7 +49,6 @@ def test_train_ctlearn_model(reco_task, dl1_gamma_file, dl1_proton_file, tmp_pat "--DLImageReader.focal_length_choice=EQUIVALENT", f"--DLImageReader.allowed_tels={allowed_tels}", ] - # Include background only for classification task if reco_task == "type": argv.extend( @@ -49,13 +58,13 @@ def test_train_ctlearn_model(reco_task, dl1_gamma_file, dl1_proton_file, tmp_pat "--DLImageReader.enforce_subarray_equality=False", ] ) - - # Run training - assert run_tool(TrainCTLearnModel(), 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" @@ -76,3 +85,16 @@ def test_train_ctlearn_model(reco_task, dl1_gamma_file, dl1_proton_file, tmp_pat f"'val_loss' values out of range [0.0, 1.0] for {reco_task}: " f"{val_loss.tolist()}" ) + # Check that the event file for TensorBoard is created for train and validation + for subfolder in ["train", "validation"]: + subfolder_path = output_dir / subfolder + assert subfolder_path.is_dir(), f"TensorBoard '{subfolder}' directory missing in {output_dir}" + # Check that at least one file starting with 'events.out.tfevents.' exists + event_files = [ + f for f in subfolder_path.iterdir() + if f.is_file() and f.name.startswith("events.out.tfevents.") + ] + assert event_files, ( + f"No TensorBoard event file starting with 'events.out.tfevents.' " + f"found in {subfolder_path}" + ) \ No newline at end of file diff --git a/ctlearn/tools/train_model.py b/ctlearn/tools/train_model.py index e4f729d2..85929d90 100644 --- a/ctlearn/tools/train_model.py +++ b/ctlearn/tools/train_model.py @@ -1,12 +1,12 @@ """ -Tool to train a ``CTLearnModel`` on R1/DL1a data using the ``DLDataReader`` and ``DLDataLoader``. +Base tool to train a ``CTLearnModel``on R1/DL1a data using the ``DLDataReader`` and ``DLDataLoader``. """ -import atexit -import keras -import pandas as pd +__all__ = ["TrainCTLearnModel"] + +from abc import abstractmethod import numpy as np -import tensorflow as tf + from ctapipe.core import Tool from ctapipe.core.tool import ToolConfigurationError @@ -22,67 +22,18 @@ ComponentName, Unicode, ) -from dl1_data_handler.reader import DLDataReader from ctlearn import __version__ as ctlearn_version -from ctlearn.core.loader import DLDataLoader from ctlearn.core.model import CTLearnModel -from ctlearn.utils import validate_trait_dict +from ctlearn.tools.utils import validate_trait_dict +from dl1_data_handler.reader import DLDataReader class TrainCTLearnModel(Tool): """ - Tool to train a ``~ctlearn.core.model.CTLearnModel`` on R1/DL1a data. - - The tool trains a CTLearn 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 - """ + Base tool to train a ``~ctlearn.core.model.CTLearnModel`` on R1/DL1a data. - name = "ctlearn-train-model" - description = __doc__ - - examples = """ - To train a CTLearn model for the classification of the primary particle type: - > ctlearn-train-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 model for the regression of the primary particle energy: - > ctlearn-train-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 model for the regression of the primary particle - arrival direction based on the offsets in camera coordinates: - > ctlearn-train-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 model for the regression of the primary particle - arrival direction based on the offsets in sky coordinates: - > ctlearn-train-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 \\ + The tool holds configurations and set up functions. + """ input_dir_signal = Path( @@ -136,7 +87,15 @@ class TrainCTLearnModel(Tool): ), ).tag(config=True) - model_type = ComponentName(CTLearnModel, default_value="ResNet").tag(config=True) + model_type = CaselessStrEnum( + ["SingleCNN", "ResNet", "LoadedModel"], + default_value="ResNet", + allow_none=False, + help=( + "Model type to be used in the Keras or PyTorch framework. " + "The framework is determined by the inherited tools being used." + ), + ).tag(config=True) output_dir = Path( exits=False, @@ -178,12 +137,6 @@ class TrainCTLearnModel(Tool): max=0.99, ).tag(config=True) - save_best_validation_only = Bool( - default_value=True, - allow_none=False, - help="Set whether to save the best validation checkpoint only.", - ).tag(config=True) - optimizer = Dict( default_value={ "name": "Adam", @@ -196,20 +149,6 @@ class TrainCTLearnModel(Tool): ), ).tag(config=True) - lr_reducing = Dict( - default_value={ - "factor": 0.5, - "patience": 5, - "min_delta": 0.01, - "min_lr": 0.000001, - }, - allow_none=True, - help=( - "Learning rate reducing parameters for the Keras callback. " - "E.g. {'factor': 0.5, 'patience': 5, 'min_delta': 0.01, 'min_lr': 0.000001}. " - ), - ).tag(config=True) - random_seed = Int( default_value=0, help=( @@ -219,19 +158,28 @@ class TrainCTLearnModel(Tool): ), ).tag(config=True) - save_onnx = Bool( - default_value=False, + save_best_validation_only = Bool( + default_value=True, allow_none=False, - help="Set whether to save model in an ONNX file.", + help="Set whether to save the best validation checkpoint only.", + ).tag(config=True) + + lr_reducing = Dict( + default_value={"factor": 0.5, "patience": 5, "min_delta": 0.01, "min_lr": 0.000001}, + allow_none=True, + help=( + "Learning rate reducing parameters for the Keras callback or the PyTorch scheduler. " + "E.g. {'factor': 0.5, 'patience': 5, 'min_delta': 0.01, 'min_lr': 0.000001}. " + ) ).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}. " - ), + help=( + "Early stopping parameters for the Keras callback or the PyTorch scheduler. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + ) ).tag(config=True) aliases = { @@ -240,6 +188,10 @@ class TrainCTLearnModel(Tool): "pattern-signal": "TrainCTLearnModel.file_pattern_signal", "pattern-background": "TrainCTLearnModel.file_pattern_background", "reco": "TrainCTLearnModel.reco_tasks", + "n-epochs": "TrainCTLearnModel.n_epochs", + "batch-size": "TrainCTLearnModel.batch_size", + "random-seed": "TrainCTLearnModel.random_seed", + "save-best-val": "TrainCTLearnModel.save_best_validation_only", ("o", "output"): "TrainCTLearnModel.output_dir", } @@ -252,10 +204,6 @@ def setup(self): raise ToolConfigurationError( f"Output directory {self.output_dir} already exists." ) - # 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) # Get signal input files self.input_url_signal = [] for signal_pattern in self.file_pattern_signal: @@ -321,235 +269,40 @@ def setup(self): ) # Set up the data loaders for training and validation - indices = list(range(self.dl1dh_reader._get_n_events())) + 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(indices) - n_validation_examples = int( + np.random.shuffle(self.indices) + self.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] - self.training_loader = DLDataLoader( - self.dl1dh_reader, - training_indices, - tasks=self.reco_tasks, - 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( - self.dl1dh_reader, - validation_indices, - tasks=self.reco_tasks, - 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.training_indices = self.indices[self.n_validation_examples:] + self.validation_indices = self.indices[:self.n_validation_examples] - # Set up the callbacks - monitor = "val_loss" - monitor_mode = "min" - # Model checkpoint callback - model_path = f"{self.output_dir}/ctlearn_model.keras" - model_checkpoint_callback = keras.callbacks.ModelCheckpoint( - filepath=model_path, - monitor=monitor, - verbose=1, - mode=monitor_mode, - save_best_only=self.save_best_validation_only, - ) - # Tensorboard callback - tensorboard_callback = keras.callbacks.TensorBoard( - log_dir=self.output_dir, histogram_freq=1 - ) - # CSV logger callback - 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, - ] + # Validate the optimizer parameters + validate_trait_dict(self.optimizer, ["name", "base_learning_rate"]) + self.learning_rate = self.optimizer["base_learning_rate"] + self.adam_epsilon = self.optimizer.get("adam_epsilon", 1e-8) + # Validate the learning rate reducing parameters + if self.lr_reducing is not None: + validate_trait_dict( + self.lr_reducing, ["factor", "patience", "min_delta", "min_lr"] + ) + # Validate the early stopping parameters if self.early_stopping is not None: - # EarlyStopping callback 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"], - ) - 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"] - ) - lr_reducing_callback = keras.callbacks.ReduceLROnPlateau( - monitor=monitor, - factor=self.lr_reducing["factor"], - patience=self.lr_reducing["patience"], - mode=monitor_mode, - verbose=1, - min_delta=self.lr_reducing["min_delta"], - min_lr=self.lr_reducing["min_lr"], - ) - self.callbacks.append(lr_reducing_callback) - - def start(self): - - # Open a strategy scope. - with self.strategy.scope(): - # Construct the model - self.log.info("Setting up the model.") - self.model = CTLearnModel.from_name( - self.model_type, - input_shape=self.training_loader.input_shape, - tasks=self.reco_tasks, - parent=self, - ).model - # 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"] - # Set the epsilon for the Adam optimizer - adam_epsilon = None - if self.optimizer["name"] == "Adam": - # Validate the epsilon for the Adam optimizer - validate_trait_dict(self.optimizer, ["adam_epsilon"]) - # Set the epsilon for the Adam optimizer - adam_epsilon = self.optimizer["adam_epsilon"] - # Select optimizer with appropriate arguments - # Dict of optimizer_name: (optimizer_fn, optimizer_args) - optimizers = { - "Adadelta": ( - keras.optimizers.Adadelta, - dict(learning_rate=learning_rate), - ), - "Adam": ( - keras.optimizers.Adam, - dict(learning_rate=learning_rate, epsilon=adam_epsilon), - ), - "RMSProp": ( - keras.optimizers.RMSprop, - dict(learning_rate=learning_rate), - ), - "SGD": (keras.optimizers.SGD, dict(learning_rate=learning_rate)), - } - # Get the optimizer function and arguments - optimizer_fn, optimizer_args = optimizers[self.optimizer["name"]] - # Get the losses and metrics for the model - 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 - ) - - # Train and evaluate the model - self.log.info("Training and evaluating...") - self.model.fit( - self.training_loader, - validation_data=self.validation_loader, - epochs=self.n_epochs, - class_weight=self.dl1dh_reader.class_weight, - callbacks=self.callbacks, - verbose=2, - ) - self.log.info("Training and evaluating finished succesfully!") + # 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 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. - - Function to build the fully connected head of the CTLearn model using the specified parameters. - - Parameters - ---------- - inputs : keras.layers.Layer - Keras layer of the model. - layers : dict - Dictionary containing the number of neurons (as value) in the fully connected head for each task (as key). - tasks : list - List of tasks to build the head for. - - Returns - ------- - logits : dict - Dictionary containing the logits for each task. - """ - losses, metrics = {}, {} - if "type" in self.reco_tasks: - losses["type"] = keras.losses.CategoricalCrossentropy( - reduction="sum_over_batch_size" - ) - metrics["type"] = [ - keras.metrics.CategoricalAccuracy(name="accuracy"), - keras.metrics.AUC(name="auc"), - ] - # 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(tasks) == 1: - losses = losses["type"] - metrics = metrics["type"] - if "energy" in self.reco_tasks: - losses["energy"] = keras.losses.MeanAbsoluteError( - reduction="sum_over_batch_size" - ) - metrics["energy"] = keras.metrics.MeanAbsoluteError(name="mae_energy") - if "cameradirection" in self.reco_tasks: - losses["cameradirection"] = keras.losses.MeanAbsoluteError( - reduction="sum_over_batch_size" - ) - 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 - - -def main(): - # Run the tool - tool = TrainCTLearnModel() - tool.run() - - -if __name__ == "main": - main() diff --git a/ctlearn/tools/utils.py b/ctlearn/tools/utils.py new file mode 100644 index 00000000..0ed8425a --- /dev/null +++ b/ctlearn/tools/utils.py @@ -0,0 +1,216 @@ +""" +Utility functions for the CTLearn tools. +""" + +from enum import Enum +import pathlib +from importlib.resources import files, as_file +import os +import time +from tqdm import tqdm + +import tensorflow as tf +import torch + +from ctapipe.core import Provenance +from ctapipe.core.tool import ToolConfigurationError +from ctapipe.core.traits import TraitError +from ctapipe.instrument.optics import FocalLengthKind +from ctapipe.instrument import SubarrayDescription + + +__all__ = [ + "monitor_progress", + "validate_trait_dict", + "get_lst1_subarray_description", + "FrameworkType", + "setup_framework", +] + +def monitor_progress(src_path, dst_path, stop_event, logger): + try: + total_size = os.path.getsize(src_path) + except OSError: + logger.error(f"Unable to access source file '{src_path}'.") + return + + last_logged_percent = -1 + + with tqdm(total=total_size, unit='B', unit_scale=True, desc="Copy Progress") as pbar: + while not stop_event.is_set(): + try: + current_size = os.path.getsize(dst_path) + except OSError: + current_size = 0 + + pbar.n = current_size + pbar.refresh() + + # Logging cada 10% + if total_size > 0: + percent = int((current_size / total_size) * 100) + if percent // 10 != last_logged_percent // 10: + logger.info(f"Progress: {percent}%") + last_logged_percent = percent + + time.sleep(0.5) + # Ensure the progress bar reaches the end + try: + final_size = os.path.getsize(dst_path) + pbar.n = final_size + pbar.refresh() + logger.info("Copy completed.") + except OSError: + logger.warning("Could not get final size of output file.") + +def validate_trait_dict(dict, required_keys): + """ + Validate that a dictionary contains all required keys. + + Parameters + ---------- + dict : dict + Dictionary to validate. + required_keys : set + Set of required keys. + + Returns + ------- + bool + True if the dictionary contains all required keys. Otherwise, raises a TraitError. + """ + missing_keys = required_keys - dict.keys() + if missing_keys: + raise TraitError(f"Dict is missing required key(s): {', '.join(missing_keys)}") + return True + +def get_lst1_subarray_description(focal_length_choice=FocalLengthKind.EFFECTIVE): + """ + Load subarray description from bundled file + + Parameters + ---------- + focal_length_choice : FocalLengthKind + Choice of focal length to use. Options are ``FocalLengthKind.EQUIVALENT`` + and ``FocalLengthKind.EFFECTIVE``. Default is ``FocalLengthKind.EFFECTIVE``. + + Returns + ------- + SubarrayDescription + Subarray description of the LST-1 telescope. + """ + with as_file(files("ctlearn") / "resources/LST-1_SubarrayDescription.h5") as path: + Provenance().add_input_file(path, role="SubarrayDescription") + return SubarrayDescription.from_hdf(path, focal_length_choice=focal_length_choice) + +class FrameworkType(Enum): + """ + Deep learning framework type enumeration. + + This enumeration specifies which deep learning framework to use for + model training and inference. CTLearn supports both Keras (TensorFlow backend) + and PyTorch frameworks. + + Attributes: + KERAS (int): Use Keras/TensorFlow framework (value: 1) + - Advantages: High-level API, easy to use, good for prototyping + - TensorFlow 2.x with Keras API + - Suitable for production deployment + + PYTORCH (int): Use PyTorch framework (value: 2) + - Advantages: Dynamic computation graphs, flexible, research-friendly + - PyTorch 1.x or 2.x + - Better for custom architectures and experimental models + + Example: + >>> from ctlearn.core.ctlearn_enum import FrameworkType + >>> framework = FrameworkType.PYTORCH + >>> print(framework.name) # 'PYTORCH' + >>> print(framework.value) # 'PyTorch' + """ + KERAS = "Keras" + PYTORCH = "PyTorch" + + +def setup_framework(model_paths): + """ + Detects the deep learning framework from model paths, ensures consistency, + and configures the hardware devices for distributed or single-device execution. + + This function iterates through the provided model paths, infers the framework + (Keras or PyTorch) based on file extensions, and ensures all models belong + to the same framework. It then initializes the appropriate hardware setup + (e.g., MirroredStrategy for Keras, CUDA/CPU device for PyTorch). + + Parameters + ---------- + model_paths : list of str or pathlib.Path or None + A list of paths pointing to the saved model files. `None` values are ignored. + + Returns + ------- + tuple + A tuple containing: + - framework_type (FrameworkType): The identified framework enum (FrameworkType.KERAS or FrameworkType.PYTORCH). + - num_devices (int): The number of devices available/configured for the framework. + - strategy (tf.distribute.Strategy or None): The TensorFlow distribution strategy if Keras is detected, else None. + - device (torch.device or None): The PyTorch target device if PyTorch is detected, else None. + + Raises + ------ + ToolConfigurationError + If no valid model paths are provided, or if multiple inconsistent frameworks + are detected across the provided paths. + """ + + def _detect_framework(path_val): + """ + Determines framework based on file extension. + Returns 'Keras', 'PyTorch', or raises TraitError. + """ + path = pathlib.Path(path_val) + ext = path.suffix.lower() + if ext in [".keras", ".h5"]: + return FrameworkType["KERAS"] + elif ext in [".pt", ".pth"]: + return FrameworkType["PYTORCH"] + else: + raise TraitError( + f"Invalid model extension '{ext}' for file '{path}'. " + "Expected '.keras' or '.h5' for Keras, or '.pt' or '.pth' for PyTorch." + ) + + # Detect frameworks from all non-None paths + detected_frameworks = {} + for path in model_paths: + if path is not None: + detected_frameworks[path] = _detect_framework(path) + # Fail immediately if no valid model paths are provided + if not detected_frameworks: + raise ToolConfigurationError( + "No model paths were specified. At least one valid model path " + "(.keras/.h5 for Keras, or .pt/.pth for PyTorch) must be provided." + ) + # Verify consistency across specified paths + unique_frameworks = set(detected_frameworks.values()) + if len(unique_frameworks) > 1: + details = ", ".join(f"{p}: {fw.value}" for p, fw in detected_frameworks.items()) + raise ToolConfigurationError( + f"Inconsistent model frameworks detected across paths: {details}. " + "All specified model files must belong to the same framework." + ) + # Set framework directly from the detected FrameworkType enum + framework_type = next(iter(unique_frameworks)) + # Configure framework-specific hardware/device setup + strategy, device = None, None + if framework_type == FrameworkType.KERAS: + strategy = tf.distribute.MirroredStrategy() + num_devices = strategy.num_replicas_in_sync + elif framework_type == FrameworkType.PYTORCH: + if torch.cuda.is_available(): + device = torch.device("cuda") + num_devices = torch.cuda.device_count() + else: + device = torch.device("cpu") + num_devices = 1 + return framework_type, num_devices, strategy, device \ No newline at end of file diff --git a/ctlearn/utils.py b/ctlearn/utils.py deleted file mode 100644 index d2929527..00000000 --- a/ctlearn/utils.py +++ /dev/null @@ -1,49 +0,0 @@ -from importlib.resources import files, as_file - -from ctapipe.core import Provenance -from ctapipe.core.traits import TraitError -from ctapipe.instrument import SubarrayDescription -from ctapipe.instrument.optics import FocalLengthKind - - -__all__ = ["get_lst1_subarray_description", "validate_trait_dict"] - -def get_lst1_subarray_description(focal_length_choice=FocalLengthKind.EFFECTIVE): - """ - Load subarray description from bundled file - - Parameters - ---------- - focal_length_choice : FocalLengthKind - Choice of focal length to use. Options are ``FocalLengthKind.EQUIVALENT`` - and ``FocalLengthKind.EFFECTIVE``. Default is ``FocalLengthKind.EFFECTIVE``. - - Returns - ------- - SubarrayDescription - Subarray description of the LST-1 telescope. - """ - with as_file(files("ctlearn") / "resources/LST-1_SubarrayDescription.h5") as path: - Provenance().add_input_file(path, role="SubarrayDescription") - return SubarrayDescription.from_hdf(path, focal_length_choice=focal_length_choice) - -def validate_trait_dict(dict, required_keys): - """ - Validate that a dictionary contains all required keys. - - Parameters - ---------- - dict : dict - Dictionary to validate. - required_keys : set - Set of required keys. - - Returns - ------- - bool - True if the dictionary contains all required keys. Otherwise, raises a TraitError. - """ - missing_keys = required_keys - dict.keys() - if missing_keys: - raise TraitError(f"Dict is missing required key(s): {', '.join(missing_keys)}") - return True diff --git a/docs/source/conf.py b/docs/source/conf.py index a7e93629..83edd800 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -16,6 +16,9 @@ import sys import ctlearn +# Point Sphinx to the root directory where `ctlearn` lives +sys.path.insert(0, os.path.abspath('../..')) + # -- Project information ----------------------------------------------------- project = 'CTLearn' diff --git a/docs/source/usage.rst b/docs/source/usage.rst index c49050e1..c411bf61 100644 --- a/docs/source/usage.rst +++ b/docs/source/usage.rst @@ -12,11 +12,12 @@ This page provides a brief overview of how to use the CTLearn tools. Training tool ------------- -To train a model, use the `ctlearn-train-model` command. The following command will display all available options for training a CTLearn model: +To train a model, use the `ctlearn-train-keras-model` or `ctlearn-train-pytorch-model` command. The following command will display all available options for training a CTLearn model: .. code-block:: bash - ctlearn-train-model --help-all + ctlearn-train-keras-model --help-all + ctlearn-train-pytorch-model --help-all View training progress in real time with TensorBoard: @@ -27,7 +28,7 @@ View training progress in real time with TensorBoard: Prediction tools ---------------- -To predict with a trained model, use the `ctlearn-predict-mono-model` or `ctlearn-predict-stereo-model` command. The following command will display all available options for predicting with a CTLearn model: +To predict with a trained Keras or PyTorch model, use the `ctlearn-predict-mono-model` or `ctlearn-predict-stereo-model` command. The following command will display all available options for predicting with a CTLearn model: .. code-block:: bash diff --git a/pyproject.toml b/pyproject.toml index f2a5c057..7c662447 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,19 +28,24 @@ classifiers = [ requires-python = ">=3.12" dependencies = [ - "dl1_data_handler>=0.14.8", "astropy", + "ctapipe[all]>=0.29", + "dl1_data_handler>=0.14.8", + "numba", "numpy", "pandas", "pip", + "pydot", + "pytorch-lightning", "pyyaml", "scikit-learn", - "numba", - "tensorflow>=2.16", - "tensorboard", - "pydot", "setuptools", - "ctapipe[all]>=0.29", + "tensorboard", + "tensorflow>=2.16", + "torch>=2.4.0", + "torchvision", + "torchmetrics", + "opencv-python", ] dynamic = ["version"] @@ -49,19 +54,41 @@ dynamic = ["version"] packages = ["ctlearn"] [project.optional-dependencies] -doc = [ + +# all is with all optional *runtime* dependencies +# use `dev` to get really all dependencies +all = [ + "matplotlib ~=3.0", + "pyirf ~=0.14.0", +] + +tests = [ + # at the moment, essentially all tests rely on test data from simtel + # it doesn't make sense to skip all of these. + "ctlearn", + "pytest >=9.0", + "pytest-cov", + "pytest-xdist", + "pytest_astropy_header", +] + +docs = [ "sphinx", "sphinx-rtd-theme", ] -# self reference allows all to be defined in terms of other extras -all = ["ctlearn[doc]"] + +dev = [ + "ctlearn[all,docs,tests]", + "setuptools_scm[toml]", +] [project.urls] 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-train-pytorch-model = "ctlearn.tools.pytorch.train_model:main" ctlearn-predict-mono-model = "ctlearn.tools.predict_model:mono_tool" ctlearn-predict-stereo-model = "ctlearn.tools.predict_model:stereo_tool" ctlearn-predict-LST1= "ctlearn.tools.predict_LST1:main"