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..00f01395 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.2.*'] exclude: - python-version: '3.13' tensorflow-version: '2.16.*' - python-version: '3.14' tensorflow-version: '2.16.*' + - python-version: '3.13' + torch-version: '2.2.*' + - python-version: '3.14' + torch-version: '2.2.*' 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/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..6e808a07 100644 --- a/README.rst +++ b/README.rst @@ -47,6 +47,49 @@ 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. +Running CTLearn Training and Prediction +--------------------------------------- + +CTLearn provides a unified command line interface (CLI) using `ctapipe`'s ``Tool`` and ``Component`` systems, supporting both the **Keras** and **PyTorch** deep learning frameworks. + +Launching training +~~~~~~~~~~~~~~~~~~ + +You can launch a training run using the unified tool ``ctlearn-train``. To run with a specific framework, set the ``--framework`` option (choices: ``keras``, ``pytorch``): + +.. code-block:: bash + + # Launch training with PyTorch + ctlearn-train --framework=pytorch --output ./my_output_dir --signal /path/to/signal/h5/ --pattern-signal "*.dl1.h5" --reco energy --n_epochs=10 --batch_size=32 + + # Launch training with Keras + ctlearn-train --framework=keras --output ./my_output_dir --signal /path/to/signal/h5/ --pattern-signal "*.dl1.h5" --reco energy --n_epochs=10 --batch_size=32 + +Common Training Command Options: +* ``--framework``: Deep learning framework to use (``keras`` or ``pytorch``). +* ``-o``, ``--output``: Directory to save experiment checkpoints, parameters, and logs. +* ``--signal``: Directory containing signal HDF5 data files. +* ``--pattern-signal``: File name pattern for signal files (e.g. ``*.dl1.h5``). +* ``--reco``: Tasks to train (e.g. ``type``, ``energy``, ``cameradirection``, ``skydirection``). Multiple tasks can be provided. +* ``--n_epochs``: Number of epochs to train. +* ``--batch_size``: Batch size for training. +* ``--save_onnx=True``: Export the trained model to ONNX format. +* ``--load_onnx_model=PATH``: Load an existing ONNX model to train/fine-tune. +* ``--overwrite``: Overwrite the output directory if it already exists. + +Launching prediction +~~~~~~~~~~~~~~~~~~~~ + +Similarly, predictions on test/observation data can be executed using the unified prediction tools (for monoscopic or stereoscopic mode): + +.. code-block:: bash + + # Monoscopic prediction with PyTorch + ctlearn-predict-mono-model --framework=pytorch --output ./pred_results --signal /path/to/test/h5/ --pattern-signal "*.dl1.h5" --energy_checkpoint /path/to/checkpoint.pth + + # Monoscopic prediction with Keras + ctlearn-predict-mono-model --framework=keras --output ./pred_results --signal /path/to/test/h5/ --pattern-signal "*.dl1.h5" --energy_checkpoint /path/to/keras_model/ + Citing this software -------------------- diff --git a/condaenv.vowayxih.requirements.txt b/condaenv.vowayxih.requirements.txt new file mode 100644 index 00000000..b29ef163 --- /dev/null +++ b/condaenv.vowayxih.requirements.txt @@ -0,0 +1,11 @@ +numba +tensorflow>=2.14,<2.15 +dl1_data_handler>=0.14.1,<0.15 +pydot +pytorch_lightning +deepspeed +onnx +onnxsim +torch==2.5.0 +torchvision==0.20.0 +torchaudio==2.5.0 \ No newline at end of file diff --git a/count_types.py b/count_types.py new file mode 100644 index 00000000..d2e19910 --- /dev/null +++ b/count_types.py @@ -0,0 +1,231 @@ +from ctlearn.core.data_loader.loader import DLDataLoader +from dl1_data_handler.reader import DLImageReader +from dl1_data_handler.reader import DLDataReader +import numpy as np +import matplotlib.pyplot as plt +from ctlearn.tools.train.pytorch.utils import read_configuration +import os +def on_key(event): + # Check if the "Esc" key was pressed + if event.key == 'escape': + plt.close(event.canvas.figure) + exit() + + + +gamma_dir = "/storage/ctlearn_data/h5_files/mc/gamma-diffuse/" + +gamma_list= [ +"gamma_theta_16.087_az_108.090_runs123-182.dl1.h5", +# "gamma_theta_16.087_az_108.090_runs1-62.dl1.h5", +# "gamma_theta_16.087_az_108.090_runs183-242.dl1.h5", + +"gamma_theta_16.087_az_251.910_runs121-180.dl1.h5", +# "gamma_theta_16.087_az_251.910_runs1-60.dl1.h5", +# "gamma_theta_16.087_az_251.910_runs181-240.dl1.h5", + +"gamma_theta_23.161_az_260.739_runs129-187.dl1.h5", +# "gamma_theta_23.161_az_260.739_runs188-246.dl1.h5", +# "gamma_theta_23.161_az_260.739_runs247-305.dl1.h5", + +"gamma_theta_23.161_az_99.261_runs118-176.dl1.h5", +# "gamma_theta_23.161_az_99.261_runs1-59.dl1.h5", +# "gamma_theta_23.161_az_99.261_runs177-235.dl1.h5", + +"gamma_theta_30.390_az_266.360_runs121-180.dl1.h5", +# "gamma_theta_30.390_az_266.360_runs1-60.dl1.h5", +# "gamma_theta_30.390_az_266.360_runs181-240.dl1.h5", + +"gamma_theta_30.390_az_93.640_runs121-180.dl1.h5", +"gamma_theta_30.390_az_93.640_runs1-60.dl1.h5", +# "gamma_theta_30.390_az_93.640_runs181-240.dl1.h5", + +"gamma_theta_37.661_az_270.641_runs121-180.dl1.h5", +# "gamma_theta_37.661_az_270.641_runs1-60.dl1.h5", +# "gamma_theta_37.661_az_270.641_runs181-240.dl1.h5", +# "gamma_theta_37.661_az_270.641_runs241-300.dl1.h5", + +"gamma_theta_37.661_az_89.359_runs121-180.dl1.h5", +# "gamma_theta_37.661_az_89.359_runs1-60.dl1.h5", +# "gamma_theta_37.661_az_89.359_runs181-240.dl1.h5", +# "gamma_theta_37.661_az_89.359_runs241-300.dl1.h5", + +"gamma_theta_6.000_az_180.000_runs121-180.dl1.h5", +# "gamma_theta_6.000_az_180.000_runs1-60.dl1.h5", +# "gamma_theta_6.000_az_180.000_runs181-240.dl1.h5", +# "gamma_theta_6.000_az_180.000_runs241-300.dl1.h5", + +"gamma_theta_9.579_az_126.888_runs121-180.dl1.h5", +# "gamma_theta_9.579_az_126.888_runs1-60.dl1.h5", +# "gamma_theta_9.579_az_126.888_runs181-240.dl1.h5", +# "gamma_theta_9.579_az_126.888_runs241-300.dl1.h5", + +"gamma_theta_9.579_az_233.112_runs121-180.dl1.h5", +# "gamma_theta_9.579_az_233.112_runs1-60.dl1.h5", +# "gamma_theta_9.579_az_233.112_runs181-240.dl1.h5", +# "gamma_theta_9.579_az_233.112_runs241-300.dl1.h5", +] + + +proton_dir = "/storage/ctlearn_data/h5_files/mc/protons/" + +proton_list=[ +"proton_theta_16.087_az_108.090_runs1-416.dl1.h5", +# "proton_theta_16.087_az_108.090_runs417-834.dl1.h5", +# "proton_theta_16.087_az_108.090_runs835-1250.dl1.h5", +"proton_theta_16.087_az_251.910_runs1-417.dl1.h5", +# "proton_theta_16.087_az_251.910_runs418-834.dl1.h5", +# "proton_theta_16.087_az_251.910_runs835-1250.dl1.h5", +"proton_theta_23.161_az_260.739_runs1-417.dl1.h5", +# "proton_theta_23.161_az_260.739_runs418-834.dl1.h5", +# "proton_theta_23.161_az_260.739_runs835-1250.dl1.h5", +"proton_theta_23.161_az_99.261_runs1-417.dl1.h5", +# "proton_theta_23.161_az_99.261_runs418-832.dl1.h5", +# "proton_theta_23.161_az_99.261_runs833-1250.dl1.h5", +"proton_theta_30.390_az_266.360_runs1-416.dl1.h5", +# "proton_theta_30.390_az_266.360_runs417-834.dl1.h5", +# "proton_theta_30.390_az_266.360_runs835-1250.dl1.h5", +"proton_theta_30.390_az_93.640_runs1-420.dl1.h5", +# "proton_theta_30.390_az_93.640_runs421-834.dl1.h5", +# "proton_theta_30.390_az_93.640_runs835-1250.dl1.h5", +"proton_theta_37.661_az_270.641_runs1-421.dl1.h5", +# "proton_theta_37.661_az_270.641_runs422-836.dl1.h5", +# "proton_theta_37.661_az_270.641_runs837-1250.dl1.h5", +"proton_theta_37.661_az_89.359_runs1-406.dl1.h5", +# "proton_theta_37.661_az_89.359_runs408-867.dl1.h5", +# "proton_theta_37.661_az_89.359_runs868-1250.dl1.h5", +"proton_theta_6.000_az_180.000_runs1-416.dl1.h5", +# "proton_theta_6.000_az_180.000_runs417-830.dl1.h5", +# "proton_theta_6.000_az_180.000_runs831-1250.dl1.h5", +"proton_theta_9.579_az_126.888_runs1-417.dl1.h5", +# "proton_theta_9.579_az_126.888_runs418-834.dl1.h5", +# "proton_theta_9.579_az_126.888_runs835-1250.dl1.h5", + +"proton_theta_9.579_az_233.112_runs1-417.dl1.h5", +# "proton_theta_9.579_az_233.112_runs418-834.dl1.h5", +# "proton_theta_9.579_az_233.112_runs835-1250.dl1.h5", +] + + +config_file = "./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml" + + +# dl1dh_reader = DLDataReader.from_name( +# "DLImageReader", +# input_url_signal=[file], +# channels = ["cleaned_image","cleaned_peak_time"], +# # input_url_background=sorted(self.input_url_background), +# # parent=self, +# ) + +cnt_type=0 + +for file_name in gamma_list: + dl1dh_reader = DLDataReader.from_name( + "DLImageReader", + input_url_signal=[os.path.join(gamma_dir,file_name)], + channels = ["cleaned_image","cleaned_peak_time"], + # input_url_background=sorted(self.input_url_background), + # parent=self, + ) + print(f"file: {file_name} events: {dl1dh_reader.n_signal_events}") + cnt_type+=dl1dh_reader.n_signal_events + + +# for file_name in proton_list: +# dl1dh_reader = DLDataReader.from_name( +# "DLImageReader", +# input_url_signal=[os.path.join(proton_dir,file_name)], +# channels = ["cleaned_image","cleaned_peak_time"], +# # input_url_background=sorted(self.input_url_background), +# # parent=self, +# ) +# print(f"file: {file_name} events: {dl1dh_reader.n_signal_events}") +# cnt_type+=dl1dh_reader.n_signal_events + +print("cnt_type: ", cnt_type) + + +gamma_cnt = "3389398" +proton_cnt = "3590228" + + +gamma_cnt = 1234488 +proton_cnt = 1249970 + +# # dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) +# # dataloader = DLDataLoader.create("pytorch") +# parameters = read_configuration(config_file) + +# random_seed = 0 +# indices = list(range(dl1dh_reader._get_n_events())) +# training_loader = DLDataLoader.create( +# framework="pytorch", +# DLDataReader=dl1dh_reader, +# indices=indices, +# tasks=["type","energy","skydirection","cameradirection","hillas"], +# batch_size=32, +# random_seed=0, +# sort_by_intensity=False, +# stack_telescope_images=False, +# parameters= parameters, +# use_augmentation=True +# ) + + +# print(len(training_loader)) +# ii=0 +# for batch_idx, (features, labels) in enumerate(training_loader): + +# plt.rcParams['keymap.quit'].append(' ') +# fig, axes = plt.subplots(1, 2, figsize=(15, 5)) +# ax = axes.ravel() +# # fig.canvas.mpl_connect('key_press_event', lambda evt: print(repr(evt.key))) +# fig.canvas.mpl_connect('key_press_event', on_key) + +# if len(features)>0: +# for id in range(len(features["image"])): +# # gammaness = 0 +# # gammaness = labels["gammaness"][id] + +# # if gammaness>0.9: + +# image = features["image"][id] +# # clean_image = features["clean_image"] +# peak_time = features["peak_time"][id] +# # clean_peak_time = features["clean_peak_time"] +# labels_class = labels["type"][id] + +# hillas_intensity = features["hillas"]["hillas_intensity"][id] +# leakage_pixels_width_1= features["hillas"]["leakage_pixels_width_1"][id] +# leakage_pixels_width_2= features["hillas"]["leakage_pixels_width_2"][id] + +# leakage_intensity_width_1 = features["hillas"]["leakage_intensity_width_1"][id] +# leakage_intensity_width_2 = features["hillas"]["leakage_intensity_width_2"][id] + +# morphology_n_islands= features["hillas"]["morphology_n_islands"][id] +# image = np.transpose(image, (1, 2, 0)) +# # clean_image = np.transpose(clean_image, (1, 2, 0)) + +# peak_time = np.transpose(peak_time, (1, 2, 0)) +# # clean_peak_time = np.transpose(clean_peak_time, (1, 2, 0)) + +# ax[0].set_title( +# f"Charge:\n Hillas Intensity:{hillas_intensity} \n Leakage_p_w_2: {leakage_pixels_width_2} \n Leakage_i_w_2: {leakage_intensity_width_2} \n morphology_n_islands: {morphology_n_islands} \n labels_class: {str(labels_class)}" , fontsize=8 +# ) +# # ax[1].set_title( +# # f"Peak time: \n Gammaness: {gammaness}" , fontsize=8 +# # ) +# ax[0].imshow(image, cmap="viridis") +# ax[1].imshow(peak_time, cmap="viridis") + +# # print(f"Gammaness: {gammaness}") +# # if leakage_intensity_width_2<0.2: +# # print(f"found {leakage_intensity_width_2}") +# plt.show() +# # plt.show(block=False) +# # plt.pause(1) +# # print(batch_idx) +# plt.close() + +# ii = 0 \ No newline at end of file diff --git a/ctlearn/__init__.py b/ctlearn/__init__.py index c4c2735c..a6291244 100644 --- a/ctlearn/__init__.py +++ b/ctlearn/__init__.py @@ -1,3 +1,23 @@ +import sys +import os +import warnings + +# Suppress noisy logs globally before any other imports occur +is_debug = '--debug' in sys.argv or any(arg.startswith('--log-level=DEBUG') for arg in sys.argv) +if not is_debug: + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' + os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' + os.environ['NCCL_DEBUG'] = 'WARN' + warnings.filterwarnings("ignore", category=UserWarning) + warnings.filterwarnings("ignore", ".*NoneDefaultNotAllowedWarning.*") + warnings.filterwarnings("ignore", ".*MergeConflictWarning.*") + warnings.filterwarnings("ignore", ".*'ctlearn.tools.train_model' found in sys.modules.*") + from ._version import __version__ +import importlib.util + + +def is_package_available(package_name: str) -> bool: + return importlib.util.find_spec(package_name) is not None -__all__ = ["__version__"] +__all__ = ["__version__", "is_package_available"] diff --git a/ctlearn/all_tools.svg b/ctlearn/all_tools.svg new file mode 100644 index 00000000..0ac41b42 --- /dev/null +++ b/ctlearn/all_tools.svg @@ -0,0 +1 @@ +ctlearncoremodelloadertoolsbase_train_modelctlearn_enumkeras.train_keras_modelpredict_LST1predict_modelpytorch.train_pytorch_modeltrain_modelCTLearnModelattention: Noneattention: NoneLoadedModelmodel: Nonebackbone_model: Noneinput_layer: Nonelogits: Nonemodel: NoneDLDataLoaderDLDataReader: Noneindices: Nonetasks: Nonebatch_size: Nonerandom_seed: Nonestack_telescope_images: Nonesort_by_intensity: Noneinput_shape: Noneinput_shape: Noneinput_shape: NoneTrainCTLearnModelFrameworkTypeKERAS: 1PYTORCH: 2TqdmProgressBarTrainKerasModelLST1PredictionToolMonoPredictCTLearnModelPredictCTLearnModelStereoPredictCTLearnModelTrainPyTorchModelDLFrameWorkGenerated bypy2puml \ No newline at end of file diff --git a/ctlearn/conftest.py b/ctlearn/conftest.py index cdfcdb18..fa5e8733 100644 --- a/ctlearn/conftest.py +++ b/ctlearn/conftest.py @@ -7,6 +7,7 @@ import numpy as np import pytest import shutil +from unittest import mock from astropy import units as u from astropy.table import Column, Table from traitlets.config.loader import Config @@ -14,7 +15,7 @@ from ctapipe.core import run_tool from ctapipe.io import write_table from ctapipe.utils import get_dataset_path -from ctlearn.tools import TrainCTLearnModel +from ctlearn.tools import DLFrameWork from ctlearn.utils import get_lst1_subarray_description @@ -240,41 +241,41 @@ def ctlearn_trained_r1_mono_models(r1_gamma_file, r1_proton_file, tmp_path_facto telescope_type = "LST" # Loop over reconstruction tasks and train models for each combination ctlearn_trained_r1_mono_models = {} - for reco_task in ["type", "energy", "cameradirection"]: - # Output directory for trained model - output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" - - # Build command-line arguments - argv = [ - f"--signal={signal_dir}", - "--pattern-signal=*.r1.h5", - f"--output={output_dir}", - f"--reco={reco_task}", - "--TrainCTLearnModel.n_epochs=1", - "--TrainCTLearnModel.batch_size=2", - "--TrainCTLearnModel.dl1dh_reader_type=DLWaveformReader", - "--DLWaveformReader.sequence_length=5", - "--DLWaveformReader.focal_length_choice=EQUIVALENT", - ] - - # Include background only for classification task - if reco_task == "type": - argv.extend( - [ - f"--background={background_dir}", - "--pattern-background=*.r1.h5", - "--DLWaveformReader.enforce_subarray_equality=False", - ] - ) - - # Run training - assert run_tool(TrainCTLearnModel(config=config), argv=argv, cwd=tmp_path) == 0 + with mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True): + for reco_task in ["type", "energy", "cameradirection"]: + # Output directory for trained model + output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" + + # Build command-line arguments + argv = [ + f"--signal={signal_dir}", + "--pattern-signal=*.r1.h5", + f"--output={output_dir}", + f"--reco={reco_task}", + "--TrainCTLearnModel.n_epochs=1", + "--TrainCTLearnModel.batch_size=2", + "--TrainCTLearnModel.dl1dh_reader_type=DLWaveformReader", + "--DLWaveformReader.sequence_length=5", + "--DLWaveformReader.focal_length_choice=EQUIVALENT", + ] + + # Include background only for classification task + if reco_task == "type": + argv.extend( + [ + f"--background={background_dir}", + "--pattern-background=*.r1.h5", + ] + ) + + # Run training + assert run_tool(DLFrameWork(config=config), argv=argv, cwd=tmp_path) == 0 - 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() + 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() return ctlearn_trained_r1_mono_models @@ -322,46 +323,46 @@ def ctlearn_trained_dl1_mono_models(dl1_gamma_file, dl1_proton_file, tmp_path_fa # Loop over telescope types and reconstruction tasks # and train models for each combination ctlearn_trained_dl1_mono_models = {} - for telescope_type, allowed_tels in telescope_types.items(): - for reco_task in ["type", "energy", "cameradirection"]: - # Output directory for trained model - output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" - - # Build command-line arguments - argv = [ - f"--signal={signal_dir}", - "--pattern-signal=*.dl1.h5", - f"--output={output_dir}", - f"--reco={reco_task}", - "--TrainCTLearnModel.n_epochs=1", - "--TrainCTLearnModel.batch_size=2", - "--DLImageReader.focal_length_choice=EQUIVALENT", - f"--DLImageReader.allowed_tels={allowed_tels}", - ] - - # Include background only for classification task - if reco_task == "type": - argv.extend( - [ - f"--background={background_dir}", - "--pattern-background=*.dl1.h5", - "--DLImageReader.enforce_subarray_equality=False", - f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", - ] + with mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True): + for telescope_type, allowed_tels in telescope_types.items(): + for reco_task in ["type", "energy", "cameradirection"]: + # Output directory for trained model + output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" + + # Build command-line arguments + argv = [ + f"--signal={signal_dir}", + "--pattern-signal=*.dl1.h5", + f"--output={output_dir}", + f"--reco={reco_task}", + "--TrainCTLearnModel.n_epochs=1", + "--TrainCTLearnModel.batch_size=2", + "--DLImageReader.focal_length_choice=EQUIVALENT", + f"--DLImageReader.allowed_tels={allowed_tels}", + ] + + # Include background only for classification task + if reco_task == "type": + argv.extend( + [ + f"--background={background_dir}", + "--pattern-background=*.dl1.h5", + f"--DLImageReader.image_mapper_type={image_mapper_types[telescope_type]}", + ] + ) + + # Run training + assert ( + run_tool(DLFrameWork(config=config), argv=argv, cwd=tmp_path) == 0 ) - # 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() + 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 @@ -403,42 +404,42 @@ 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", - "--TrainCTLearnModel.stack_telescope_images=True", - "--DLImageReader.mode=stereo", - "--DLImageReader.focal_length_choice=EQUIVALENT", - f"--DLImageReader.allowed_tels={allowed_tels}", - ] - - # Include background only for classification task - if reco_task == "type": - argv.extend( - [ - f"--background={background_dir}", - "--pattern-background=*.dl1.h5", - "--DLImageReader.enforce_subarray_equality=False", - ] - ) - - # Run training - assert run_tool(TrainCTLearnModel(config=config), argv=argv, cwd=tmp_path) == 0 + with mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True): + for reco_task in ["type", "energy", "skydirection"]: + # Output directory for trained model + output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}" + + # Build command-line arguments + argv = [ + f"--signal={signal_dir}", + "--pattern-signal=*.dl1.h5", + f"--output={output_dir}", + f"--reco={reco_task}", + "--TrainCTLearnModel.n_epochs=1", + "--TrainCTLearnModel.batch_size=2", + "--TrainCTLearnModel.stack_telescope_images=True", + "--DLImageReader.mode=stereo", + "--DLImageReader.focal_length_choice=EQUIVALENT", + f"--DLImageReader.allowed_tels={allowed_tels}", + ] + + # Include background only for classification task + if reco_task == "type": + argv.extend( + [ + f"--background={background_dir}", + "--pattern-background=*.dl1.h5", + ] + ) + + # Run training + assert run_tool(DLFrameWork(config=config), argv=argv, cwd=tmp_path) == 0 - 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() + 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 diff --git a/ctlearn/core/__init__.py b/ctlearn/core/__init__.py new file mode 100644 index 00000000..996469a1 --- /dev/null +++ b/ctlearn/core/__init__.py @@ -0,0 +1,2 @@ +"""ctlearn command line tools. +""" diff --git a/ctlearn/core/ctlearn_enum.py b/ctlearn/core/ctlearn_enum.py new file mode 100644 index 00000000..506abd57 --- /dev/null +++ b/ctlearn/core/ctlearn_enum.py @@ -0,0 +1,239 @@ +""" +CTLearn Enumeration Types Module + +This module defines enumeration types used throughout the CTLearn framework +to ensure type safety and consistency when specifying framework types, tasks, +event types, and operation modes. + +Enumerations: + FrameworkType: Deep learning framework selection (Keras or PyTorch) + Task: Machine learning task type (classification, regression, etc.) + EventType: Cosmic ray particle type classification + Mode: Operation mode for the CTLearn pipeline +""" + +from enum import Enum + +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) # 2 + """ + KERAS = 1 + PYTORCH = 2 + + +class Task(Enum): + """ + Machine learning task type enumeration. + + This enumeration defines the different analysis tasks that CTLearn can perform + on Cherenkov telescope data. Each task corresponds to a specific scientific + goal in gamma-ray astronomy. + + Attributes: + type (int): Particle type classification task (value: 0) + - Classify events as gamma-ray or background (proton/electron) + - Binary classification problem + - Output: Class probabilities (gamma vs hadron) + - Critical for gamma-ray source detection + + energy (int): Energy regression task (value: 1) + - Estimate the energy of the primary cosmic ray + - Regression problem with log-scale energy + - Output: Energy in TeV (typically log10(E/TeV)) + - Essential for measuring source spectra + + direction (int): Generic direction reconstruction task (value: 2) + - General direction estimation (deprecated, use specific types) + - Maintained for backward compatibility + + cameradirection (int): Direction reconstruction in camera coordinates (value: 3) + - Predict shower direction in camera frame + - Output: (dx, dy, distance) offsets in meters + - Must be transformed to sky coordinates for analysis + - Faster computation, no coordinate transformations needed + + skydirection (int): Direction reconstruction in sky coordinates (value: 4) + - Predict shower direction in horizontal (Alt/Az) frame + - Output: (altitude, azimuth, angular_separation) in degrees + - Directly usable for source localization + - Requires telescope pointing information + + all (int): Multi-task learning (all tasks simultaneously) (value: 5) + - Train model to perform all tasks jointly + - Shared feature extraction, task-specific heads + - Can improve performance through transfer learning + - More complex training but potentially better features + + Notes: + - Tasks can be combined in multi-task learning setups + - Each task requires specific loss functions and evaluation metrics + - The choice of task affects data preprocessing and model architecture + + Example: + >>> from ctlearn.core.ctlearn_enum import Task + >>> task = Task.energy + >>> if task == Task.energy: + ... print("Performing energy regression") + >>> print(task.name) # 'energy' + """ + type = 0 + energy = 1 + direction = 2 + cameradirection = 3 + skydirection = 4 + all = 5 + +class EventType(Enum): + """ + Cosmic ray event type enumeration. + + This enumeration classifies the primary particle that initiated the + air shower detected by the Cherenkov telescopes. Distinguishing between + gamma rays and background particles is crucial for gamma-ray astronomy. + + Attributes: + gamma (int): Gamma-ray event (value: 0) + - Primary particle: High-energy photon + - Characteristics: + * Electromagnetic shower + * Narrow, elliptical image + * Low muon content + * Preferred class for gamma-ray astronomy + - Used for source studies and spectral analysis + + proton (int): Proton-induced event (value: 1) + - Primary particle: Proton (cosmic ray) + - Characteristics: + * Hadronic shower + * Irregular, fragmented image + * High muon content + * Most common background (~90% of cosmic rays) + - Main source of background contamination + + electron (int): Electron-induced event (value: 2) + - Primary particle: Electron or positron + - Characteristics: + * Electromagnetic shower (similar to gamma) + * Can be difficult to distinguish from gamma + * Less common than protons + * Secondary background source + - Often grouped with gamma for some analyses + + Notes: + - In binary classification, typically gamma vs (proton + electron) + - Event type is known only for simulated data (Monte Carlo) + - Real observations have unknown event types (classification goal) + - Other particles (nuclei, muons) exist but are less common + + Physical Context: + - Gamma rays: Signal we want to detect + - Protons: Dominant background (~1000x more frequent) + - Electrons: Minor background component + - Background rejection crucial for sensitivity + + Example: + >>> from ctlearn.core.ctlearn_enum import EventType + >>> event = EventType.gamma + >>> if event == EventType.gamma: + ... print("Signal event detected") + >>> print(event.name) # 'gamma' + >>> print(event.value) # 0 + """ + gamma = 0 + proton = 1 + electron = 2 + +class Mode(Enum): + """ + Operation mode enumeration for the CTLearn pipeline. + + This enumeration defines the different operational modes that CTLearn + can run in, determining what actions are performed on the data. + + Attributes: + train (int): Training mode (value: 0) + - Train model on training dataset + - Update model weights through backpropagation + - Validate on validation set each epoch + - Save checkpoints and training curves + - Enables data augmentation + - Uses training-specific preprocessing + + results (int): Results generation mode (value: 1) + - Generate predictions on test or validation data + - No weight updates + - Save predictions to HDF5 files (DL2 format) + - Compute performance metrics + - Create evaluation plots and tables + - Used for final model evaluation + + validate (int): Validation mode (value: 2) + - Evaluate model on validation dataset + - No weight updates + - Compute metrics only (no predictions saved) + - Quick performance check + - Can be run during or after training + + observation (int): Real observation mode (value: 3) + - Process real telescope observations (not simulations) + - No ground truth labels available + - Generate DL2 data from DL1 real data + - Apply trained model to unknown events + - Output used for science analysis + + tunning (int): Hyperparameter tuning mode (value: 4) + - Optimize model hyperparameters + - Multiple training runs with different configurations + - Uses validation set for hyperparameter selection + - May use techniques like grid search, random search, or Bayesian optimization + - Typically automated with tools like Optuna or Ray Tune + + Typical Workflow: + 1. train: Develop and train models on simulated data + 2. validate: Quick performance checks during development + 3. tunning: Optimize hyperparameters for best performance + 4. results: Final evaluation and analysis on test set + 5. observation: Apply to real telescope data + + Notes: + - Each mode may have different data loading behavior + - Some preprocessing steps (augmentation) only active in train mode + - observation mode handles real data without labels + - Mode affects logging, checkpointing, and output format + + Example: + >>> from ctlearn.core.ctlearn_enum import Mode + >>> mode = Mode.train + >>> if mode == Mode.train: + ... print("Enabling data augmentation") + >>> elif mode == Mode.observation: + ... print("Processing real data") + >>> print(mode.name) # 'train' + """ + train = 0 + results = 1 + validate = 2 + observation = 3 + tunning = 4 diff --git a/ctlearn/core/data_loader/__init__.py b/ctlearn/core/data_loader/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/data_loader/base_loader.py b/ctlearn/core/data_loader/base_loader.py new file mode 100644 index 00000000..a850f6f2 --- /dev/null +++ b/ctlearn/core/data_loader/base_loader.py @@ -0,0 +1,320 @@ +""" +Base Data Loader Module + +This module provides an abstract base class for data loaders in CTLearn. +It defines the common interface and initialization logic for framework-specific +data loaders (Keras, PyTorch) while handling telescope data in both mono and +stereo observation modes. + +Classes: + BaseDLDataLoader: Abstract base class for deep learning data loaders +""" + +from abc import ABC, abstractmethod +import numpy as np +import random +import cv2 +from ctlearn.core.ctlearn_enum import Task + +class BaseDLDataLoader(ABC): + """ + Abstract base class for deep learning data loaders. + + This class provides the common interface and initialization logic for loading + and processing Cherenkov telescope data. It handles both mono (single telescope) + and stereo (multiple telescopes) observation modes and supports various data + processing options like sorting by intensity and stacking telescope images. + + Attributes: + DLDataReader: The data reader instance for accessing telescope event data + indices (list): List of event indices to load from the dataset + tasks (list): List of tasks to perform (e.g., classification, energy, direction) + batch_size (int): Number of batches per epoch + random_seed (int or None): Random seed for reproducibility + stack_telescope_images (bool): Whether to stack images from multiple telescopes + sort_by_intensity (bool): Whether to sort telescope images by Hillas intensity + input_shape (tuple): Shape of input data (height, width, channels) + + Methods: + __len__: Abstract method to return the number of batches per epoch + __getitem__: Abstract method to generate one batch of data + on_epoch_end: Abstract method called at the end of each epoch + """ + + def __init__( + self, + DLDataReader, + indices, + tasks, + batch_size=64, + random_seed=None, + sort_by_intensity=False, + stack_telescope_images=False, + **kwargs, + ): + """ + Initialize the base data loader. + + Sets up the data loader with a data reader, configures batch processing + parameters, and determines the input shape based on observation mode + (mono vs stereo) and image stacking options. + + Args: + DLDataReader: Instance of a data reader class (e.g., DLHDFDataReader) + that provides access to telescope event data + indices (list): List of integer indices specifying which events to load + from the dataset + tasks (list): List of Task enum values specifying which tasks to perform + (e.g., [Task.type, Task.energy, Task.cameradirection]) + batch_size (int, optional): Number of samples to include in each batch. + Defaults to 64 + random_seed (int or None, optional): Seed for random number generator + to ensure reproducibility. If None, random behavior is not deterministic. + Defaults to None + sort_by_intensity (bool, optional): If True, sort telescope images by + their Hillas intensity in descending order. Useful for stereo analysis. + Defaults to False + stack_telescope_images (bool, optional): If True, stack images from multiple + telescopes along the channel dimension. Only applicable in stereo mode. + Defaults to False + **kwargs: Additional keyword arguments passed to parent classes + """ + super().__init__() + + # Store initialization parameters + self.DLDataReader = DLDataReader + self.indices = indices + self.tasks = tasks + self.batch_size = batch_size + self.random_seed = random_seed + self.stack_telescope_images = stack_telescope_images + self.sort_by_intensity = sort_by_intensity + + # Get parent configuration if available + parent = getattr(self.DLDataReader, "parent", None) + + # Helper to get attributes with fallbacks + def get_val(name, default): + if parent is not None and hasattr(parent, name): + return getattr(parent, name) + if name in kwargs: + return kwargs[name] + if "parameters" in kwargs and isinstance(kwargs["parameters"], dict): + params = kwargs["parameters"] + if name == "use_augmentation": + return params.get("augmentation", {}).get("use_augmentation", default) + elif name == "aug_prob": + return params.get("augmentation", {}).get("aug_prob", default) + elif name == "rot_prob": + return params.get("augmentation", {}).get("rot_prob", default) + elif name == "trans_prob": + return params.get("augmentation", {}).get("trans_prob", default) + elif name == "flip_hor_prob": + return params.get("augmentation", {}).get("flip_hor_prob", default) + elif name == "flip_ver_prob": + return params.get("augmentation", {}).get("flip_ver_prob", default) + elif name == "mask_prob": + return params.get("augmentation", {}).get("mask_prob", default) + elif name == "mask_dvr_prob": + return params.get("augmentation", {}).get("mask_dvr_prob", default) + elif name == "noise_prob": + return params.get("augmentation", {}).get("noise_prob", default) + elif name == "max_rot": + return params.get("augmentation", {}).get("max_rot", default) + elif name == "max_trans": + return params.get("augmentation", {}).get("max_trans", default) + elif name == "apply_log_scaling": + return params.get("normalization", {}).get("apply_log_scaling", default) + elif name == "use_clean": + return params.get("normalization", {}).get("use_clean", default) + elif name == "use_clean_dvr": + return params.get("normalization", {}).get("use_clean_dvr", default) + elif name == "type_mu": + return params.get("normalization", {}).get("type_mu", default) + elif name == "type_sigma": + return params.get("normalization", {}).get("type_sigma", default) + elif name == "dir_mu": + return params.get("normalization", {}).get("dir_mu", default) + elif name == "dir_sigma": + return params.get("normalization", {}).get("dir_sigma", default) + elif name == "energy_mu": + return params.get("normalization", {}).get("energy_mu", default) + elif name == "energy_sigma": + return params.get("normalization", {}).get("energy_sigma", default) + elif name == "leakage_intensity_cutoff": + return params.get("cut-off", {}).get("leakage_intensity", default) + elif name == "intensity_cutoff": + return params.get("cut-off", {}).get("intensity", default) + return default + + # Initialize pre-processing & augmentation options + self.use_augmentation = get_val("use_augmentation", False) + self.aug_prob = get_val("aug_prob", 0.5) + self.rot_prob = get_val("rot_prob", 0.5) + self.trans_prob = get_val("trans_prob", 0.5) + self.flip_hor_prob = get_val("flip_hor_prob", 0.5) + self.flip_ver_prob = get_val("flip_ver_prob", 0.5) + self.mask_prob = get_val("mask_prob", 0.5) + self.mask_dvr_prob = get_val("mask_dvr_prob", 0.5) + self.noise_prob = get_val("noise_prob", 0.5) + self.max_rot = get_val("max_rot", 5.0) + self.max_trans = get_val("max_trans", 10.0) + + self.apply_log_scaling = get_val("apply_log_scaling", [True, True]) + self.use_clean = get_val("use_clean", True) + self.use_clean_dvr = get_val("use_clean_dvr", False) + + self.type_mu = get_val("type_mu", 0.0) + self.type_sigma = get_val("type_sigma", 1.0) + self.dir_mu = get_val("dir_mu", 0.0) + self.dir_sigma = get_val("dir_sigma", 1.0) + self.energy_mu = get_val("energy_mu", 0.0) + self.energy_sigma = get_val("energy_sigma", 1.0) + + self.leakage_intensity_cutoff = get_val("leakage_intensity_cutoff", 0.2) + self.intensity_cutoff = get_val("intensity_cutoff", 50.0) + + # Determine input shape based on reader type and observation mode + # Feature vector readers don't have spatial dimensions + if self.DLDataReader.__class__.__name__ != "DLFeatureVectorReader": + + # Mono mode: single telescope per event + if self.DLDataReader.mode == "mono": + # Use the input shape directly from the data reader + self.input_shape = self.DLDataReader.input_shape + + # Stereo mode: multiple telescopes per event + elif self.DLDataReader.mode == "stereo": + # Get input shape from the first selected telescope + # All telescopes are assumed to have the same image dimensions + self.input_shape = self.DLDataReader.input_shape[ + list(self.DLDataReader.selected_telescopes)[0] + ] + + # Modify input shape if stacking telescope images + # Original shape: (num_telescopes, height, width, channels) + # Stacked shape: (height, width, num_telescopes * channels) + if self.stack_telescope_images: + self.input_shape = ( + self.input_shape[1], # height + self.input_shape[2], # width + self.input_shape[0] * self.input_shape[3], # stacked channels + ) + + def clean_data(self, image, peak_time): + # Remove negative numbers and avoid inf or nans + image[image < 0] = 0 + peak_time[peak_time < 0] = 0 + image[np.isnan(image)] = 0 + image[np.isinf(image)] = 0 + peak_time[np.isnan(peak_time)] = 0 + peak_time[np.isinf(peak_time)] = 0 + return image, peak_time + + def normalize_data(self, image, peak_time, task): + # Normalization + if task == Task.type or task == "type": + image = (image - self.type_mu) / self.type_sigma + peak_time = (peak_time - self.type_mu) / self.type_sigma + elif task == Task.energy or task == "energy": + image = (image - self.energy_mu) / self.energy_sigma + peak_time = (peak_time - self.energy_mu) / self.energy_sigma + elif task in [Task.cameradirection, Task.skydirection, Task.direction, "cameradirection", "skydirection", "direction"]: + image = (image - self.dir_mu) / self.dir_sigma + peak_time = (peak_time - self.dir_mu) / self.dir_sigma + + return image, peak_time + + + def apply_augmentation(self, image, peak_time, task): + for id_batch in range(image.shape[0]): + random_aug = random.random() + + if random_aug > self.aug_prob: + if task not in [Task.cameradirection, Task.skydirection, Task.direction, "cameradirection", "skydirection", "direction"]: + random_aug_flip_ver = random.random() + if random_aug_flip_ver > self.flip_ver_prob: + image[id_batch] = np.expand_dims(cv2.flip(image[id_batch].astype(np.float32), 0), axis=-1) + peak_time[id_batch] = np.expand_dims(cv2.flip(peak_time[id_batch].astype(np.float32), 0), axis=-1) + continue + random_aug_flip_hor = random.random() + if random_aug_flip_hor > self.flip_hor_prob: + image[id_batch] = np.expand_dims(cv2.flip(image[id_batch].astype(np.float32), 1), axis=-1) + peak_time[id_batch] = np.expand_dims(cv2.flip(peak_time[id_batch].astype(np.float32), 1), axis=-1) + continue + random_aug_rot = random.random() + if random_aug_rot > self.rot_prob: + (h, w) = image[id_batch].shape[:2] + angle = random.uniform(-self.max_rot, self.max_rot) + scale = 1.0 + center = (w // 2, h // 2) + rotation_matrix = cv2.getRotationMatrix2D(center, angle, scale) + image[id_batch] = np.expand_dims(cv2.warpAffine( + image[id_batch].astype(np.float32), rotation_matrix, (w, h) + ), axis=-1) + peak_time[id_batch] = np.expand_dims(cv2.warpAffine( + peak_time[id_batch].astype(np.float32), rotation_matrix, (w, h) + ), axis=-1) + continue + random_aug_trans = random.random() + if random_aug_trans > self.trans_prob: + (h, w) = image[id_batch].shape[:2] + tx = random.uniform(-self.max_trans, self.max_trans) + ty = random.uniform(-self.max_trans, self.max_trans) + translation_matrix = np.float32([[1, 0, tx], [0, 1, ty]]) + image[id_batch] = np.expand_dims(cv2.warpAffine( + image[id_batch].astype(np.float32), translation_matrix, (w, h) + ), axis=-1) + peak_time[id_batch] = np.expand_dims(cv2.warpAffine( + peak_time[id_batch].astype(np.float32), translation_matrix, (w, h) + ), axis=-1) + continue + return image, peak_time + + @abstractmethod + def __len__(self): + """ + Get the number of batches per epoch. + + This method must be implemented by subclasses to return the total number + of batches that will be generated in one epoch, typically calculated as + ceil(total_samples / batch_size). + + Returns: + int: Number of batches per epoch + """ + pass + + @abstractmethod + def __getitem__(self, index): + """ + Generate one batch of data. + + This method must be implemented by subclasses to return a single batch + of data at the specified index. The batch should contain input features + and corresponding labels formatted appropriately for the framework. + + Args: + index (int): Index of the batch to generate (0 to len(self) - 1) + + Returns: + tuple: A tuple containing: + - features (dict or array): Input features for the batch + - labels (dict or array): Ground truth labels for the batch + - metadata (optional): Additional information about the batch + """ + pass + + @abstractmethod + def on_epoch_end(self): + """ + Perform operations at the end of each epoch. + + This method must be implemented by subclasses to perform any necessary + cleanup or updates at the end of an epoch. Common operations include + shuffling indices for the next epoch or updating internal state. + + This method is typically called automatically by the training framework + after processing all batches in an epoch. + """ + pass diff --git a/ctlearn/core/data_loader/keras_loader.py b/ctlearn/core/data_loader/keras_loader.py new file mode 100644 index 00000000..5c09ecda --- /dev/null +++ b/ctlearn/core/data_loader/keras_loader.py @@ -0,0 +1,452 @@ +""" +Keras Data Loader Module + +This module provides a Keras-specific data loader implementation for CTLearn. +It extends both Keras Sequence and the base data loader to provide efficient +batch generation for training with Keras/TensorFlow models. + +Classes: + KerasDLDataLoader: Keras Sequence for loading and preprocessing telescope data +""" + +import numpy as np +import keras +from keras.utils import Sequence, to_categorical +from .base_loader import BaseDLDataLoader + +from dl1_data_handler.reader import ProcessType + +class KerasDLDataLoader(Sequence, BaseDLDataLoader): + """ + Keras data loader for Cherenkov telescope data. + + This class implements the Keras Sequence interface to provide efficient + data loading for training. It supports both monoscopic (single telescope) + and stereoscopic (multiple telescopes) observation modes, with options + for image stacking and intensity-based sorting. + + Inherits from: + Sequence: Keras sequence for thread-safe batch generation + BaseDLDataLoader: Base class providing common data loading functionality + + Attributes: + Inherited from BaseDLDataLoader including: + - DLDataReader: Data reader for accessing telescope events + - indices: Array of event indices to process + - tasks: List of tasks (type, energy, direction) + - batch_size: Number of samples per batch + - random_seed: Seed for shuffling + - sort_by_intensity: Whether to sort by Hillas intensity + - stack_telescope_images: Whether to stack images from multiple telescopes + """ + + def __init__( + self, + **kwargs, + ): + """ + Initialize the Keras data loader. + + This constructor initializes the data loader by calling the parent + class constructors and performing initial index shuffling if a + random seed is provided. + + Args: + **kwargs: Keyword arguments passed to BaseDLDataLoader, including: + - DLDataReader: Data reader instance + - indices: Event indices to load + - tasks: List of tasks to perform + - batch_size: Batch size + - random_seed: Random seed for shuffling + - sort_by_intensity: Whether to sort by intensity + - stack_telescope_images: Whether to stack images + """ + BaseDLDataLoader.__init__(self, **kwargs) + Sequence.__init__(self) + # Perform initial shuffling of indices if random seed is set + self.on_epoch_end() + + def __len__(self): + """ + Get the number of batches per epoch. + + This method calculates the number of batches required to cover the entire dataset + based on the batch size. Uses floor division to ensure complete batches only. + + Returns: + int: Number of batches per epoch (total_samples // batch_size) + + Note: + Samples that don't fit into a complete batch are not included in the epoch. + For example, with 100 samples and batch_size=32, this returns 3 (96 samples used). + """ + return int(np.floor(len(self.indices) / self.batch_size)) + + def on_epoch_end(self): + """ + Update indices after each epoch. + + This method is called automatically by Keras at the end of each epoch. + If a random seed is provided, it shuffles the indices to provide different + batch compositions in each epoch, which can improve training convergence. + + The shuffling is deterministic (uses the same seed each time) to ensure + reproducibility while still providing epoch-to-epoch variation. + + Side Effects: + - If random_seed is not None: Shuffles self.indices in-place using + the configured random seed + - If random_seed is None: No operation performed + """ + if self.random_seed is not None: + # Set random seed for reproducibility + np.random.seed(self.random_seed) + # Shuffle indices in-place to randomize batch composition + np.random.shuffle(self.indices) + + def __getitem__(self, index): + """ + Generate one batch of data. + + This method is called by Keras to generate one batch of data based on + the provided index. It delegates to mode-specific methods (_get_mono_item + or _get_stereo_item) depending on the observation mode. + + Args: + index (int): Index of the batch to generate (0 to len(self) - 1) + + Returns: + tuple: (features, labels) where: + - features: Input data formatted for the model + * Keras 2: dict with 'input' key + * Keras 3: numpy array directly + * Shape varies by mode and configuration + - labels: Ground truth labels as dict or array + * Dict keys depend on tasks (type, energy, skydirection, cameradirection) + * For single task classification: categorical array instead of dict + + Note: + The batch indices are computed as: + batch_indices = self.indices[index * batch_size : (index + 1) * batch_size] + """ + # Generate indices of the batch + batch_indices = self.indices[ + index * self.batch_size : (index + 1) * self.batch_size + ] + features, labels = None, None + + # Generate batch based on observation mode + if self.DLDataReader.mode == "mono": + # Monoscopic mode: single telescope per event + batch = self.DLDataReader.generate_mono_batch(batch_indices) + features, labels = self._get_mono_item(batch) + elif self.DLDataReader.mode == "stereo": + # Stereoscopic mode: multiple telescopes per event + batch = self.DLDataReader.generate_stereo_batch(batch_indices) + features, labels = self._get_stereo_item(batch) + return features, labels + + def _get_mono_item(self, batch): + """ + Retrieve features and labels for one batch of monoscopic data. + + This method extracts and formats data from a single-telescope batch, + preparing features and task-specific labels for training or inference. + + Args: + batch (astropy.table.Table): Table containing monoscopic event data + Expected columns include: + - features: Telescope images or feature vectors + - true_shower_primary_class: Particle type (0=gamma, 1=proton) + - log_true_energy: Logarithm of true energy + - fov_lon, fov_lat: Sky direction in field-of-view coordinates + - cam_coord_offset_x, cam_coord_offset_y: Camera coordinate offsets + + Returns: + tuple: (features, labels) where: + - features: Input features formatted for Keras + * Keras 2: dict with 'input' key containing numpy array + * Keras 3: numpy array directly + * Shape: (batch_size, height, width, channels) + - labels: Task-specific labels + * If only 'type' task: categorical array (batch_size, 2) + * Otherwise: dict with keys matching self.tasks: + - 'type': one-hot encoded particle type + - 'energy': log energy values + - 'skydirection': (lon, lat) coordinates + - 'cameradirection': (offset_x, offset_y) coordinates + """ + # Initialize labels dictionary + labels = {} + + # Retrieve telescope images and store in features dictionary + features = {"input": batch["features"].data} + + if features["input"].shape[-1] == 2: + image = features["input"][..., 0:1] + peak_time = features["input"][..., 1:2] + + active_task = self.tasks[0] if self.tasks else None + + image, peak_time = self.clean_data(image, peak_time) + + if self.use_augmentation: + image, peak_time = self.apply_augmentation(image, peak_time, active_task) + + image, peak_time = self.normalize_data(image, peak_time, active_task) + features["input"] = np.concatenate([image, peak_time], axis=-1) + else: + # If it's a waveform or just image, just use it directly for now (or apply cleaning if needed) + pass + + + # Extract particle type classification labels + if "type" in self.tasks: + # Convert to one-hot encoding (0=gamma, 1=proton) + labels["type"] = to_categorical( + batch["true_shower_primary_class"].data, + num_classes=2, + ) + # Temporary fix: Use array instead of dict for single-task classification + # Required until Keras fully supports class weights for multiple outputs + # See: https://github.com/keras-team/keras/issues/11735 + if len(self.tasks) == 1: + labels = to_categorical( + batch["true_shower_primary_class"].data, + num_classes=2, + ) + + # Extract energy regression labels + if "energy" in self.tasks: + # Energy is already in log scale + labels["energy"] = batch["log_true_energy"].data + + # Extract sky direction reconstruction labels + if "skydirection" in self.tasks: + # Stack longitude and latitude into single array + labels["skydirection"] = np.stack( + ( + batch["fov_lon"].data, + batch["fov_lat"].data, + batch["angular_separation"].data, + ), + axis=1, + ) + + # Extract camera direction reconstruction labels + if "cameradirection" in self.tasks: + # Stack camera x and y offsets into single array + labels["cameradirection"] = np.stack( + ( + batch["cam_coord_offset_x"].data, + batch["cam_coord_offset_y"].data, + batch["cam_coord_distance"].data, + ), + axis=1, + ) + + # Temporary fix for Keras 2/3 compatibility + # Keras 3 expects arrays directly, not wrapped in dict + if int(keras.__version__.split(".")[0]) >= 3: + features = features["input"] + + return features, labels + + def _get_stereo_item(self, batch): + """ + Retrieve features and labels for one batch of stereoscopic data. + + This method processes multi-telescope events, grouping telescope data + by event, optionally sorting by intensity, and stacking images if requested. + It also handles both telescope-level and subarray-level feature vectors. + + Args: + batch (astropy.table.Table): Table containing stereoscopic event data + Expected columns include all mono columns plus: + - obs_id: Observation run identifier + - event_id: Event identifier within observation + - tel_type_id: Telescope type identifier + - hillas_intensity: For sorting (if sort_by_intensity=True) + - mono_feature_vectors: Telescope-level features (optional) + - stereo_feature_vectors: Subarray-level features (optional) + + Returns: + tuple: (features, labels) where: + - features: Input features formatted for Keras + * If stacked images: (batch_size, height, width, n_channels * n_tel) + * If unstacked: (batch_size, n_tel, height, width, n_channels) + * Feature vectors have different shapes + - labels: Task-specific labels (same structure as _get_mono_item) + + Note: + Events are grouped by (obs_id, event_id, tel_type_id) for simulations, + or by (obs_id, event_id, tel_type_id) for observations. + Labels are extracted from the first telescope in each group since they + are event-level quantities (same for all telescopes in an event). + """ + # Initialize labels dictionary + labels = {} + + # Group batch by event to collect all telescopes for each event + if self.DLDataReader.process_type == ProcessType.Simulation: + # For simulations, group by observation, event, telescope type, and particle class + batch_grouped = batch.group_by( + ["obs_id", "event_id", "tel_type_id", "true_shower_primary_class"] + ) + elif self.DLDataReader.process_type == ProcessType.Observation: + # For real observations, particle class is unknown + batch_grouped = batch.group_by(["obs_id", "event_id", "tel_type_id"]) + + # Initialize lists for collecting event-level data + features, mono_feature_vectors, stereo_feature_vectors = [], [], [] + true_shower_primary_class = [] + log_true_energy = [] + fov_lon, fov_lat, angular_separation = [], [], [] + cam_coord_offset_x, cam_coord_offset_y, cam_coord_distance = [], [], [] + + # Process each event group + for group_element in batch_grouped.groups: + # Process telescope images if available + if "features" in batch.colnames: + # Sort telescopes by Hillas intensity if requested + if self.sort_by_intensity: + # Sort in descending order (brightest first) + group_element.sort(["hillas_intensity"], reverse=True) + + # Stack telescope images for stereo analysis if requested + if self.stack_telescope_images: + # Retrieve telescope images for this event + plain_features = group_element["features"].data + # Concatenate along channel axis: (h, w, c*n_tel) + stacked_features = np.concatenate( + [plain_features[i] for i in range(plain_features.shape[0])], + axis=-1, + ) + # Append stacked images + # Shape: (height, width, n_channels * n_telescopes) + features.append(stacked_features) + else: + # Keep telescopes as separate dimension + # Shape: (n_telescopes, height, width, n_channels) + features.append(group_element["features"].data) + + # Retrieve telescope-level feature vectors if available + if "mono_feature_vectors" in batch.colnames: + mono_feature_vectors.append(group_element["mono_feature_vectors"].data) + + # Retrieve subarray-level feature vectors if available + if "stereo_feature_vectors" in batch.colnames: + stereo_feature_vectors.append( + group_element["stereo_feature_vectors"].data + ) + + # Extract event-level labels (same for all telescopes in event) + # FIXME: This won't work correctly for divergent pointing directions + # where different telescopes point in different directions + + # Particle type classification + if "type" in self.tasks: + # Use first telescope's value (same for all telescopes in event) + true_shower_primary_class.append( + group_element["true_shower_primary_class"].data[0] + ) + + # Energy regression + if "energy" in self.tasks: + log_true_energy.append(group_element["log_true_energy"].data[0]) + + # Sky direction reconstruction + if "skydirection" in self.tasks: + fov_lon.append(group_element["fov_lon"].data[0]) + fov_lat.append(group_element["fov_lat"].data[0]) + angular_separation.append(group_element["angular_separation"].data[0]) + + # Camera direction reconstruction + 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) + cam_coord_distance.append(group_element["cam_coord_distance"].data) + + # Format labels for each task + if "type" in self.tasks: + # Convert to one-hot encoding + labels["type"] = to_categorical( + np.array(true_shower_primary_class), + num_classes=2, + ) + # Temporary fix for single-task classification + if len(self.tasks) == 1: + labels = to_categorical( + np.array(true_shower_primary_class), + num_classes=2, + ) + + if "energy" in self.tasks: + labels["energy"] = np.array(log_true_energy) + + if "skydirection" in self.tasks: + # Stack longitude and latitude + labels["skydirection"] = np.stack( + ( + np.array(fov_lon), + np.array(fov_lat), + np.array(angular_separation), + ), + axis=1, + ) + + if "cameradirection" in self.tasks: + # Stack camera coordinate offsets + labels["cameradirection"] = np.stack( + ( + np.array(cam_coord_offset_x), + np.array(cam_coord_offset_y), + np.array(cam_coord_distance), + ), + axis=1, + ) + + # Format features based on available data type + if "features" in batch.colnames: + # Telescope images + features = {"input": np.array(features)} + + features_arr = features["input"] + active_task = self.tasks[0] if self.tasks else None + + # Slicing image and peak_time based on dimensionality + if len(features_arr.shape) == 5: # Unstacked mode: (batch, tel, height, width, channels) + image = features_arr[..., 0] + peak_time = features_arr[..., 1] + + image, peak_time = self.clean_data(image, peak_time) + image, peak_time = self.normalize_data(image, peak_time, active_task) + + features["input"] = np.stack([image, peak_time], axis=-1) + else: # Stacked mode: (batch, height, width, channels) + image = features_arr[..., ::2] + peak_time = features_arr[..., 1::2] + + image, peak_time = self.clean_data(image, peak_time) + image, peak_time = self.normalize_data(image, peak_time, active_task) + + # Re-stack alternating channels + stacked = [] + for i in range(image.shape[-1]): + stacked.append(image[..., i:i+1]) + stacked.append(peak_time[..., i:i+1]) + features["input"] = np.concatenate(stacked, axis=-1) + + # TODO: Add support for using both mono and stereo feature vectors simultaneously + if "mono_feature_vectors" in batch.colnames: + # Telescope-level feature vectors + features = {"input": np.array(mono_feature_vectors)} + if "stereo_feature_vectors" in batch.colnames: + # Subarray-level feature vectors + features = {"input": np.array(stereo_feature_vectors)} + + # Temporary fix for Keras 2/3 compatibility + if int(keras.__version__.split(".")[0]) >= 3: + features = features["input"] + + return features, labels diff --git a/ctlearn/core/data_loader/loader.py b/ctlearn/core/data_loader/loader.py new file mode 100644 index 00000000..285a0922 --- /dev/null +++ b/ctlearn/core/data_loader/loader.py @@ -0,0 +1,88 @@ +""" +Data Loader Factory Module + +This module provides a factory class for creating framework-specific data loaders +for CTLearn. It supports both Keras and PyTorch frameworks and handles dynamic +imports to avoid unnecessary dependencies. + +Classes: + DLDataLoader: Factory class for creating data loaders based on the specified framework +""" + +# from .keras_loader import KerasDLDataLoader +# from .pytorch_loader import PyTorchDLDataLoader + +class DLDataLoader: + """ + Factory class for creating deep learning data loaders. + + This class provides a static factory method to instantiate the appropriate + data loader based on the specified framework (Keras or PyTorch). It uses + dynamic imports to load only the required framework dependencies. + + Methods: + create: Static factory method to create framework-specific data loaders + """ + + @staticmethod + def create(framework, **kwargs): + """ + Create a data loader for the specified framework. + + This factory method instantiates the appropriate data loader class based + on the framework parameter. It dynamically imports the required loader + class to avoid loading unnecessary dependencies for unused frameworks. + + Args: + framework (str): The framework to use ('keras' or 'pytorch') + **kwargs: Additional keyword arguments passed to the data loader constructor + These may include: + - config: Configuration dictionary + - mode: Operation mode (train, validation, test, prediction) + - data_files: List of input data files + - batch_size: Number of samples per batch + - num_workers: Number of worker processes for data loading + + Returns: + KerasDLDataLoader or PyTorchDLDataLoader: The instantiated data loader + for the specified framework + + Raises: + ValueError: If the framework is not 'keras' or 'pytorch' + ImportError: If the required framework-specific loader cannot be imported + + Examples: + >>> # Create a PyTorch data loader + >>> loader = DLDataLoader.create('pytorch', config=config, mode='train') + + >>> # Create a Keras data loader + >>> loader = DLDataLoader.create('keras', config=config, mode='validation') + """ + # Initialize dataloader variable + dataloader = None + + # Create Keras data loader + if framework == "keras": + try: + # Dynamically import Keras loader to avoid unnecessary dependencies + from .keras_loader import KerasDLDataLoader + dataloader = KerasDLDataLoader(**kwargs) + except ImportError as e: + # Raise informative error if Keras dependencies are missing + raise ImportError(f"Not possible to import KerasDLDataLoader: {e}") from e + + # Create PyTorch data loader + elif framework == "pytorch": + try: + # Dynamically import PyTorch loader to avoid unnecessary dependencies + from .pytorch_loader import PyTorchDLDataLoader + dataloader = PyTorchDLDataLoader(**kwargs) + except ImportError as e: + # Raise informative error if PyTorch dependencies are missing + raise ImportError(f"Not possible to import PyTorchDLDataLoader: {e}") from e + + # Handle unsupported framework + else: + raise ValueError(f"Unsupported framework: {framework}") + + return dataloader diff --git a/ctlearn/core/data_loader/pytorch_loader.py b/ctlearn/core/data_loader/pytorch_loader.py new file mode 100644 index 00000000..89076eae --- /dev/null +++ b/ctlearn/core/data_loader/pytorch_loader.py @@ -0,0 +1,718 @@ +import torch +import numpy as np +from torch.utils.data import Dataset +from .base_loader import BaseDLDataLoader +from dl1_data_handler.reader import ProcessType +from ctlearn.core.ctlearn_enum import Task +from astropy import units as u +import random +import cv2 + +#import concurrent.futures + +from astropy.time import Time +from astropy.coordinates import AltAz, SkyCoord +from ctapipe.coordinates import CameraFrame +from astropy import units as u + +LST_EPOCH = Time("2018-10-01T00:00:00", scale="utc") + +class PyTorchDLDataLoader(Dataset, BaseDLDataLoader): + """ + PyTorch data loader for Cherenkov telescope data. + + This class implements the PyTorch Dataset interface to provide efficient + data loading with support for data augmentation, normalization, and asynchronous + prefetching. It handles both monoscopic and stereoscopic observation modes. + + Inherits from: + Dataset: PyTorch Dataset for data loading + BaseDLDataLoader: Base class providing common data loading functionality + + Attributes: + is_training (bool): Whether the loader is used for training + executor: ThreadPoolExecutor for asynchronous batch prefetching + use_augmentation (bool): Whether to apply data augmentation + use_clean (bool): Whether to use clean events filtering + use_clean_dvr (bool): Whether to use DVR-based event filtering + task (Task): The task type (classification, energy, direction) + Various augmentation probability parameters + Various normalization parameters (mean and std) + hillas_names (list): List of Hillas parameter names to extract + """ + + def __init__( + self, + tasks, + parameters, + use_augmentation, + T=1, + is_training=False, + **kwargs, + ): + self.is_training = is_training + + super().__init__( + tasks=tasks, + parameters=parameters, + use_augmentation=use_augmentation, + **kwargs, + ) + + self.task = tasks + self.on_epoch_end() + self.set_T(T) + + self.hillas_names = [ + "obs_id", + "event_id", + "tel_id", + "hillas_intensity", + "hillas_skewness", + "hillas_kurtosis", + "hillas_fov_lon", + "hillas_fov_lat", + "hillas_r", + "hillas_phi", + "hillas_length", + "hillas_length_uncertainty", + "hillas_width", + "hillas_width_uncertainty", + "hillas_psi", + "timing_intercept", + "timing_deviation", + "timing_slope", + "leakage_pixels_width_1", + "leakage_pixels_width_2", + "leakage_intensity_width_1", + "leakage_intensity_width_2", + "concentration_cog", + "concentration_core", + "concentration_pixel", + "morphology_n_pixels", + "morphology_n_islands", + "morphology_n_small_islands", + "morphology_n_medium_islands", + "morphology_n_large_islands", + "intensity_max", + "intensity_min", + "intensity_mean", + "intensity_std", + "intensity_skewness", + "intensity_kurtosis", + "peak_time_max", + "peak_time_min", + "peak_time_mean", + "peak_time_std", + "peak_time_skewness", + "peak_time_kurtosis", + "core_psi", + ] + + def set_T(self,T): + self.T=T + + self.indices = np.tile(self.indices, self.T) + + def __len__(self): + """ + Returns the number of batches per epoch. + + This method calculates the number of batches required to cover the entire dataset + based on the batch size. + + Returns: + -------- + int + Number of batches per epoch. + """ + return int(np.floor(len(self.indices) / self.batch_size)) + + + def on_epoch_end(self): + """ + Updates indices after each epoch. If a random seed is provided, the indices are shuffled. + + This method is called at the end of each epoch to ensure that the data is shuffled + if the shuffle attribute is set to True. This helps in improving the training process + by providing the model with a different order of data in each epoch. + """ + if self.random_seed is not None: + np.random.seed(self.random_seed) + np.random.shuffle(self.indices) + + + + def _fetch_batch(self, index): + batch_indices = self.indices[index * self.batch_size : (index + 1) * self.batch_size] + if len(batch_indices) == 0: + raise IndexError(f"No data for batch index {index} (batch_indices empty)") + + if self.DLDataReader.mode == "mono": + batch = self.DLDataReader.generate_mono_batch(batch_indices) + features, labels = self._get_mono_item(batch) + elif self.DLDataReader.mode == "stereo": + batch = self.DLDataReader.generate_stereo_batch(batch_indices) + features, labels = self._get_stereo_item(batch) + return features, labels + + # Fetching batches + def __getitem__(self, index): + """ + Returns a prefetched batch for faster data loading. + + Parameters + ---------- + index : int + Index of the batch. + + Returns + ------- + features : dict + Dictionary containing input tensors. + labels : dict + Dictionary containing labels for each task. + t : int + Virtual batch index for tracking. + """ + + # Virtual batch index + t = index // int(np.ceil(len(self.indices) / self.T / self.batch_size)) + + features, labels = self._fetch_batch(index) + return features, labels, t + + # def __getitem__(self, index): + # """ + # Generate one batch of data and retrieve the features and labels. + + # This method is called to generate one batch of monoscopic and stereoscopic data based on + # the index provided. It calls either _get_mono_item(batch) or _get_stereo_item(batch) + # based on the mode of the DLDataReader. + + # Parameters: + # ----------- + # index : int + # Index of the batch to generate. + + # Returns: + # -------- + # tuple + # A tuple containing the input data as features and the corresponding labels. + # """ + + # # data_idx = index + # # data_idx = index % int((self.total_len/self.T)/self.batch_size) + # t = index // int(np.ceil(len(self.indices)/self.T/self.batch_size)) + # # Generate indices of the batch + # batch_indices = self.indices[ + # index * self.batch_size : (index + 1) * self.batch_size + # ] + # features, labels = None, None + # if self.DLDataReader.mode == "mono": + # batch = self.DLDataReader.generate_mono_batch(batch_indices) + # features, labels = self._get_mono_item(batch) + # elif self.DLDataReader.mode == "stereo": + # batch = self.DLDataReader.generate_stereo_batch(batch_indices) + # features, labels = self._get_stereo_item(batch) + + # return features, labels, t + + # def cam_to_alt_az( + # self, tel_id, focal_length, pix_rotation, tel_az, tel_alt, cam_x, cam_y + # ): + # """ + # Transform camera coordinate offsets (cam_x, cam_y) into Alt/Az sky coordinates. + + # This method converts the given camera coordinates for each telescope into sky coordinates + # (Altitude and Azimuth), using the known pointing of each telescope and camera geometry + # such as focal length and pixel rotation. + + # Parameters + # ---------- + # tel_id : list or array-like + # List of telescope IDs corresponding to each event or observation. + + # focal_length : list or array-like + # Focal length of the telescopes in meters. + + # pix_rotation : list or array-like + # Pixel rotation angles (in degrees) for each telescope camera. + + # tel_az : list or array-like + # Azimuth of telescope pointing (in radians). + + # tel_alt : list or array-like + # Altitude of telescope pointing (in radians). + + # cam_x : list or array-like + # Camera x-coordinate positions (in meters). + + # cam_y : list or array-like + # Camera y-coordinate positions (in meters). + + # Returns + # ------- + # sky_coords_alt : list + # List of reconstructed Altitude coordinates (in degrees). + + # sky_coords_az : list + # List of reconstructed Azimuth coordinates (in degrees). + # """ + # from astropy.time import Time + + # LST_EPOCH = Time("2018-10-01T00:00:00", scale="utc") + # from astropy.coordinates import AltAz, SkyCoord + # from ctapipe.coordinates import CameraFrame + # from astropy import units as u + + # # # Get telescope ground frame position + # tel_ground_frame = self.DLDataReader.subarray.tel_coords[ + # self.DLDataReader.subarray.tel_ids_to_indices(tel_id) + # ] + + # # AltAz frame setup + # altaz = AltAz( + # location=tel_ground_frame.to_earth_location(), + # obstime=LST_EPOCH, + # ) + + # # Telescope pointing SkyCoord + # fix_tel_pointing = SkyCoord( + # az=tel_az * u.rad, + # alt=tel_alt * u.rad, + # frame=altaz, + # ) + + # sky_coords_alt = [] + # sky_coords_az = [] + + # for id in range(len(focal_length)): + + # camera_frame = CameraFrame( + # focal_length=focal_length[id] * u.m, + # rotation=pix_rotation[id] * u.deg, + # telescope_pointing=fix_tel_pointing[id], + # ) + + # cam_coord = SkyCoord( + # x=cam_x[id] * u.m, y=cam_y[id] * u.m, frame=camera_frame + # ) + + # sky_coord = cam_coord.transform_to(altaz[id]) + + # sky_coords_alt.append(sky_coord.alt.to_value(u.deg).item()) + # sky_coords_az.append(sky_coord.az.to_value(u.deg).item()) + + # return sky_coords_alt, sky_coords_az + def cam_to_alt_az(self, tel_id, focal_length, pix_rotation, tel_az, tel_alt, cam_x, cam_y): + sky_coords_alt = [] + sky_coords_az = [] + + for id in range(len(focal_length)): + # Telescopio correspondiente + tel_ground_frame = self.DLDataReader.subarray.tel_coords[ + self.DLDataReader.subarray.tel_ids_to_indices(tel_id[id]) + ] + + # Frame AltAz particular para cada telescopio + altaz_frame = AltAz( + location=tel_ground_frame.to_earth_location(), + obstime=LST_EPOCH, + ) + + fix_tel_pointing = SkyCoord( + az=tel_az[id] * u.rad, + alt=tel_alt[id] * u.rad, + frame=altaz_frame, + ) + + camera_frame = CameraFrame( + focal_length=focal_length[id] * u.m, + rotation=pix_rotation[id] * u.deg, + telescope_pointing=fix_tel_pointing, + ) + + cam_coord = SkyCoord( + x=cam_x[id] * u.m, y=cam_y[id] * u.m, frame=camera_frame + ) + + # Transformar correctamente + sky_coord = cam_coord.transform_to(altaz_frame) + + sky_coords_alt.append(sky_coord.alt.to_value(u.deg).item()) + sky_coords_az.append(sky_coord.az.to_value(u.deg).item()) + + return sky_coords_alt, sky_coords_az + + def _get_mono_item(self, batch): + """ + Retrieve the features and labels for one batch of monoscopic data. + + This method is called to retrieve the features and labels for one batch of + monoscopic data. The labels are set up based on the tasks specified. + + Parameters: + ----------- + batch : astropy.table.Table + A table containing the data for the batch. + + Returns: + -------- + tuple + A tuple containing the input data as features and the corresponding labels. + """ + # Retrieve the telescope images and store in the features dictionary + labels = {} + features = {"input": batch["features"].data} + + if "type" in self.tasks: + labels["type"] = np.stack(batch["true_shower_primary_class"].data) + + labels["energy"] = batch["log_true_energy"].data + + if "skydirection" in self.tasks: + labels["skydirection"] = np.stack( + ( + batch["fov_lon"].data, + batch["fov_lat"].data, + batch["angular_separation"].data, + ), + axis=1, + ) + if "cameradirection" in self.tasks: + labels["cameradirection"] = np.stack( + ( + batch["cam_coord_offset_x"].data, + batch["cam_coord_offset_y"].data, + batch["cam_coord_distance"].data, + ), + axis=1, + ) + + if "skydirection" in labels.keys(): + labels["direction"] = labels["skydirection"] + + if "cameradirection" in labels.keys(): + labels["direction"] = labels["cameradirection"] + + # features["hillas"] = self.DLDataReader.get_parameters(batch, self.hillas_names) + features["hillas"] = self.DLDataReader.get_parameters(batch) + + image = features["input"][..., 0:1] + peak_time = features["input"][..., 1:2] + + active_task = self.tasks[0] if self.tasks else None + + image, peak_time = self.clean_data(image, peak_time) + + features_out = {} + features_out["image"] = image + features_out["peak_time"] = peak_time + for key in features["hillas"].keys(): + features["hillas"][key] = ( + torch.from_numpy(np.array(features["hillas"][key])) + .contiguous() + .unsqueeze(-1) + ) + + if "cameradirection" in self.tasks: + + tel_ids = batch["tel_id"].data + + # tel_ground_frame = self.DLDataReader.subarray.tel_coords[ + # self.DLDataReader.subarray.tel_ids_to_indices(tel_ids) + # ] + + focal_lengths = [ + self.DLDataReader.subarray.tel[ + tel_id + ].camera.geometry.frame.focal_length + for tel_id in tel_ids + ] + pix_rotations = [ + self.DLDataReader.pix_rotation[tel_id] for tel_id in tel_ids + ] + + labels["focal_length"] = np.array( + [focal.to_value(u.m) for focal in focal_lengths] + ) + labels["pix_rotation"] = np.array( + [rot.to_value(u.deg) for rot in pix_rotations] + ) + # labels["tel_ground"] = tel_ground_frame + labels["tel_ids"] = tel_ids + labels["true_alt"] = np.array([val for val in batch["true_alt"]]) + labels["true_az"] = np.array([val for val in batch["true_az"]]) + labels["tel_az"] = batch["telescope_pointing_azimuth"].data + labels["tel_alt"] = batch["telescope_pointing_altitude"].data + + # cam_x = labels["cameradirection"][:,0].cpu().numpy().squeeze(-1) + # cam_y = labels["cameradirection"][:,1].cpu().numpy().squeeze(-1) + + # sky_coords_alt, sky_coords_az = self.cam_to_alt_az(labels["tel_ids"], labels["focal_length"], labels["pix_rotation"],labels["tel_az"],labels["tel_alt"], cam_x, cam_y) + + + + if self.is_training: + N = 4 # Repeating the number of high energies + #features_out["hillas"] = features["hillas"] + #------------------------------------------------- + if self.use_augmentation: + energy_log = torch.pow(10,torch.tensor(labels["energy"].reshape(-1))) # shape [N] + high_energy_mask = energy_log > 1 # log10(E/TeV) > 0 => E > 1 TeV + + idx_to_duplicate = torch.where(high_energy_mask)[0] + + if len(idx_to_duplicate) > 0: + def duplicate_tensor(t,idx_to_duplicate): + if isinstance(t, torch.Tensor): + extra = torch.cat([t[idx_to_duplicate] for _ in range(N)], dim=0) + return torch.cat([t, extra], dim=0).contiguous() + elif isinstance(t, np.ndarray): + # Si t es 1D, t[idx_to_duplicate] ya es de shape (M,), sólo hace falta stackear + # extra = np.tile(t[idx_to_duplicate], N) + # return np.concatenate([t, extra], axis=0) + idx_to_duplicate = idx_to_duplicate.cpu().numpy() if hasattr(idx_to_duplicate, "cpu") else idx_to_duplicate + # extra = np.tile(t[idx_to_duplicate], (N, 1, 1, 1)) # si shape es (n, x, y, z) + # Mejor: stack y luego reshape + extra = np.concatenate([t[idx_to_duplicate] for _ in range(N)], axis=0) + # O si es 1D, puedes hacer + # extra = np.tile(t[idx_to_duplicate], N) + return np.concatenate([t, extra], axis=0) + + else: + raise TypeError(f"Data type not supported: {type(t)}") + + # Duplica todas las features principales + for key in features_out: + if isinstance(features_out[key], dict): + # Por ejemplo, hillas es un dict de tensores + for k in features_out[key]: + features_out[key][k] = duplicate_tensor(features_out[key][k],idx_to_duplicate) + else: + features_out[key] = duplicate_tensor(features_out[key],idx_to_duplicate) + + # Duplica las labels + for key in labels: + labels[key] = duplicate_tensor(labels[key],idx_to_duplicate) + + + if self.use_augmentation: + if isinstance(features_out["image"], torch.Tensor): + features_out["image"] = features_out["image"].cpu().numpy() + if isinstance(features_out["peak_time"], torch.Tensor): + features_out["peak_time"] = features_out["peak_time"].cpu().numpy() + + image, peak_time = self.apply_augmentation(features_out["image"], features_out["peak_time"], active_task) + else: + image, peak_time = features_out["image"], features_out["peak_time"] + + # Normalize data + image, peak_time = self.normalize_data(image, peak_time, active_task) + + # Change to channel first + image = np.transpose(image, (0, 3, 1, 2)) + peak_time = np.transpose(peak_time, (0, 3, 1, 2)) + + if len(peak_time.shape) > 0 and peak_time.size > 0: + image = np.concatenate([image, peak_time], axis=1) + + features_out["image"] = torch.from_numpy(image.copy()).contiguous().float() + if "peak_time" in features_out: + del features_out["peak_time"] + + #Create keep_idx based on configurable leakage and intensity cutoffs + hillas = features["hillas"] + leakage = np.array(hillas["leakage_intensity_width_2"]) + intensity = np.array(hillas["hillas_intensity"]) + keep_idx = np.where((leakage <= self.leakage_intensity_cutoff) & (intensity >= self.intensity_cutoff))[0] + + if not self.is_training: + keep_idx = np.arange(len(intensity)) + + # Filter features_out + for key in features_out: + features_out[key] = features_out[key][keep_idx] + + features_out["hillas"] = features["hillas"] + + for key in features["hillas"]: + features_out["hillas"][key] = (features["hillas"][key])[keep_idx] + + # Filter labels (since it's a dict too) + for key in labels: + labels[key] = labels[key][keep_idx] + + for key in labels.keys(): + labels[key] = torch.from_numpy(labels[key]).contiguous() + + if key != "type": + labels[key] = labels[key].unsqueeze(-1) + + return features_out, labels + + def _get_stereo_item(self, batch): + """ + Retrieve features and labels for one batch of stereoscopic data. + + This method processes multi-telescope events, grouping telescope data + by event, optionally sorting by intensity, and stacking images if requested. + It also handles both telescope-level and subarray-level feature vectors. + + Args: + batch (astropy.table.Table): Table containing stereoscopic event data + Expected columns include all mono columns plus: + - obs_id: Observation run identifier + - event_id: Event identifier within observation + - tel_type_id: Telescope type identifier + - hillas_intensity: For sorting (if sort_by_intensity=True) + - mono_feature_vectors: Telescope-level features (optional) + - stereo_feature_vectors: Subarray-level features (optional) + + Returns: + tuple: (features_out, labels) where: + - features_out (dict): Contains 'image' and 'peak_time' tensors + - labels (dict): Contains task-specific labels + + Note: + Events are grouped by (obs_id, event_id, tel_type_id, true_shower_primary_class) + for simulations, or by (obs_id, event_id, tel_type_id) for observations. + Labels are extracted from the first telescope in each group since they + are event-level quantities (same for all telescopes in an event). + + TODO: + This method needs full adaptation for PyTorch tensors and proper + handling of stereoscopic image stacking similar to _get_mono_item. + """ + labels = {} + if self.DLDataReader.process_type == ProcessType.Simulation: + batch_grouped = batch.group_by( + ["obs_id", "event_id", "tel_type_id", "true_shower_primary_class"] + ) + elif self.DLDataReader.process_type == ProcessType.Observation: + batch_grouped = batch.group_by(["obs_id", "event_id", "tel_type_id"]) + features, mono_feature_vectors, stereo_feature_vectors = [], [], [] + true_shower_primary_class = [] + log_true_energy = [] + fov_lon, fov_lat, angular_separation = [], [], [] + cam_coord_offset_x, cam_coord_offset_y, cam_coord_distance = [], [], [] + for group_element in batch_grouped.groups: + if "features" in batch.colnames: + if self.sort_by_intensity: + # Sort images by the hillas intensity in a given batch if requested + group_element.sort(["hillas_intensity"], reverse=True) + # Stack the telescope images for stereo analysis + if self.stack_telescope_images: + # Retrieve the telescope images + plain_features = group_element["features"].data + # Stack the telescope images along the last axis + stacked_features = np.concatenate( + [plain_features[i] for i in range(plain_features.shape[0])], + axis=-1, + ) + # Append the stacked images to the features list + # shape: (batch_size, image_shape, image_shape, n_channels * n_tel) + features.append(stacked_features) + else: + # Append the plain images to the features list + # shape: (batch_size, n_tel, image_shape, image_shape, n_channels) + features.append(group_element["features"].data) + # Retrieve the feature vectors + if "mono_feature_vectors" in batch.colnames: + mono_feature_vectors.append(group_element["mono_feature_vectors"].data) + if "stereo_feature_vectors" in batch.colnames: + stereo_feature_vectors.append( + group_element["stereo_feature_vectors"].data + ) + # Retrieve the labels for the tasks + # FIXME: This won't work for divergent pointing directions + if "type" in self.tasks: + true_shower_primary_class.append( + group_element["true_shower_primary_class"].data[0] + ) + if "energy" in self.tasks: + log_true_energy.append(group_element["log_true_energy"].data[0]) + if "skydirection" in self.tasks: + fov_lon.append(group_element["fov_lon"].data[0]) + fov_lat.append(group_element["fov_lat"].data[0]) + angular_separation.append(group_element["angular_separation"].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) + cam_coord_distance.append(group_element["cam_coord_distance"].data) + # Store the labels in the labels dictionary + if "type" in self.tasks: + labels["type"] = np.array(true_shower_primary_class) + + # Temp fix till keras support class weights for multiple outputs or I wrote custom loss + # https://github.com/keras-team/keras/issues/11735 + if len(self.tasks) == 1: + labels = np.array(true_shower_primary_class) + + if "energy" in self.tasks: + labels["energy"] = np.array(log_true_energy) + if "skydirection" in self.tasks: + labels["skydirection"] = np.stack( + ( + np.array(fov_lon), + np.array(fov_lat), + np.array(angular_separation), + ), + axis=1, + ) + if "cameradirection" in self.tasks: + labels["cameradirection"] = np.stack( + ( + np.array(cam_coord_offset_x), + np.array(cam_coord_offset_y), + np.array(cam_coord_distance), + ), + axis=1, + ) + # Store the features in the features dictionary + if "features" in batch.colnames: + features = {"input": np.array(features)} + # TODO: Add support for both feature vectors + if "mono_feature_vectors" in batch.colnames: + features = {"input": np.array(mono_feature_vectors)} + if "stereo_feature_vectors" in batch.colnames: + features = {"input": np.array(stereo_feature_vectors)} + + # Extract features array from dict + features_arr = features["input"] + active_task = self.tasks[0] if self.tasks else None + + # Slicing image and peak_time based on dimensionality + if len(features_arr.shape) == 5: # Unstacked mode: (batch, tel, height, width, channels) + image = features_arr[..., 0] + peak_time = features_arr[..., 1] + + image, peak_time = self.clean_data(image, peak_time) + image, peak_time = self.normalize_data(image, peak_time, active_task) + + image = np.transpose(image, (0, 1, 4, 2, 3)) if len(image.shape) == 5 else np.expand_dims(image, axis=2) + peak_time = np.transpose(peak_time, (0, 1, 4, 2, 3)) if len(peak_time.shape) == 5 else np.expand_dims(peak_time, axis=2) + else: # Stacked mode: (batch, height, width, channels) + image = features_arr[..., ::2] + peak_time = features_arr[..., 1::2] + + image, peak_time = self.clean_data(image, peak_time) + image, peak_time = self.normalize_data(image, peak_time, active_task) + + image = np.transpose(image, (0, 3, 1, 2)) + + if len(peak_time.shape) > 0: + peak_time = np.transpose(peak_time, (0, 3, 1, 2)) + image = np.concatenate([image, peak_time], axis=1) + + features_out = {} + features_out["image"] = torch.from_numpy(image.copy()).contiguous().float() + + # Convert labels to PyTorch tensors + for key in labels.keys(): + labels[key] = torch.from_numpy(labels[key]).contiguous() + if key != "type": + labels[key] = labels[key].unsqueeze(-1) + + return features_out, labels diff --git a/ctlearn/core/keras/__init__.py b/ctlearn/core/keras/__init__.py new file mode 100644 index 00000000..e69de29b 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..0a133a64 100644 --- a/ctlearn/core/attention.py +++ b/ctlearn/core/keras/attention.py @@ -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..fe9c4e87 --- /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. +""" + +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, +) + +__all__ = [ + "build_fully_connect_keras_head", + "KerasSingleCNN", + "KerasResNet", + "KerasLoadedModel", +] + + +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/loader.py deleted file mode 100644 index 92fd96c2..00000000 --- a/ctlearn/core/loader.py +++ /dev/null @@ -1,309 +0,0 @@ -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): - """ - Generates batches for Keras application. - - DLDataLoader is a data loader class that inherits from ``~keras.utils.Sequence``. - It is designed to handle and load data for deep learning models in a batch-wise manner. - - Attributes: - ----------- - data_reader : DLDataReader - An instance of DLDataReader to read the input data. - indices : list - List of indices to specify the data to be loaded. - tasks : list - List of tasks to be performed on the data to properly set up the labels. - batch_size : int - Size of the batch to load the data. - random_seed : int, optional - Whether to shuffle the data after each epoch with a provided random seed. - - Methods: - -------- - __len__(): - Returns the number of batches per epoch. - on_epoch_end(): - Updates indices after each epoch if random seed is provided. - __getitem__(index): - Generates one batch of data using _get_mono_item(index) or _get_stereo_item(index). - _get_mono_item(index): - Generates one batch of monoscopic data. - _get_stereo_item(index): - Generates one batch of stereoscopic data. - """ - - def __init__( - self, - DLDataReader, - indices, - tasks, - batch_size=64, - random_seed=None, - sort_by_intensity=False, - stack_telescope_images=False, - **kwargs, - ): - super().__init__(**kwargs) - "Initialization" - self.DLDataReader = DLDataReader - self.indices = indices - self.tasks = tasks - self.batch_size = batch_size - self.random_seed = random_seed - self.on_epoch_end() - self.stack_telescope_images = stack_telescope_images - self.sort_by_intensity = sort_by_intensity - - # Set the input shape based on the mode of the DLDataReader - if self.DLDataReader.__class__.__name__ != "DLFeatureVectorReader": - if self.DLDataReader.mode == "mono": - self.input_shape = self.DLDataReader.input_shape - elif self.DLDataReader.mode == "stereo": - self.input_shape = self.DLDataReader.input_shape[ - list(self.DLDataReader.selected_telescopes)[0] - ] - # Reshape inputs into proper dimensions - # for the stereo analysis with stacked images - if self.stack_telescope_images: - self.input_shape = ( - self.input_shape[1], - self.input_shape[2], - self.input_shape[0] * self.input_shape[3], - ) - - def __len__(self): - """ - Returns the number of batches per epoch. - - This method calculates the number of batches required to cover the entire dataset - based on the batch size. - - Returns: - -------- - int - Number of batches per epoch. - """ - return int(np.floor(len(self.indices) / self.batch_size)) - - def on_epoch_end(self): - """ - Updates indices after each epoch. If a random seed is provided, the indices are shuffled. - - This method is called at the end of each epoch to ensure that the data is shuffled - if the shuffle attribute is set to True. This helps in improving the training process - by providing the model with a different order of data in each epoch. - """ - if self.random_seed is not None: - np.random.seed(self.random_seed) - np.random.shuffle(self.indices) - - def __getitem__(self, index): - """ - Generate one batch of data and retrieve the features and labels. - - This method is called to generate one batch of monoscopic and stereoscopic data based on - the index provided. It calls either _get_mono_item(batch) or _get_stereo_item(batch) - based on the mode of the DLDataReader. - - Parameters: - ----------- - index : int - Index of the batch to generate. - - Returns: - -------- - tuple - A tuple containing the input data as features and the corresponding labels. - """ - # Generate indices of the batch - batch_indices = self.indices[ - index * self.batch_size : (index + 1) * self.batch_size - ] - features, labels = None, None - if self.DLDataReader.mode == "mono": - batch = self.DLDataReader.generate_mono_batch(batch_indices) - features, labels = self._get_mono_item(batch) - elif self.DLDataReader.mode == "stereo": - batch = self.DLDataReader.generate_stereo_batch(batch_indices) - features, labels = self._get_stereo_item(batch) - return features, labels - - def _get_mono_item(self, batch): - """ - Retrieve the features and labels for one batch of monoscopic data. - - This method is called to retrieve the features and labels for one batch of - monoscopic data. The labels are set up based on the tasks specified. - - Parameters: - ----------- - batch : astropy.table.Table - A table containing the data for the batch. - - Returns: - -------- - tuple - A tuple containing the input data as features and the corresponding labels. - """ - # Retrieve the telescope images and store in the features dictionary - labels = {} - features = batch["features"].data - if "type" in self.tasks: - labels["type"] = to_categorical( - batch["true_shower_primary_class"].data, - num_classes=2, - ) - # Temp fix till keras support class weights for multiple outputs or I wrote custom loss - # https://github.com/keras-team/keras/issues/11735 - if len(self.tasks) == 1: - labels = to_categorical( - batch["true_shower_primary_class"].data, - num_classes=2, - ) - if "energy" in self.tasks: - labels["energy"] = batch["log_true_energy"].data - if "skydirection" in self.tasks: - labels["skydirection"] = np.stack( - ( - batch["fov_lon"].data, - batch["fov_lat"].data, - ), - axis=1, - ) - if "cameradirection" in self.tasks: - labels["cameradirection"] = np.stack( - ( - batch["cam_coord_offset_x"].data, - batch["cam_coord_offset_y"].data, - ), - axis=1, - ) - return features, labels - - def _get_stereo_item(self, batch): - """ - Retrieve the features and labels for one batch of stereoscopic data. - - This method is called to retrieve the features and labels for one batch of - stereoscopic data. The original batch is grouped to retrieve the telescope - data for each event and then the telescope images or waveforms are stored - by the hillas intensity or stacked if required. Feature vectors can also - be retrieved if available for ``telescope``- and ``subarray``level. The - labels are set up based on the tasks specified. - - Parameters: - ----------- - batch : astropy.table.Table - A table containing the data for the batch. - - Returns: - -------- - tuple - A tuple containing the input data as features and the corresponding labels. - """ - labels = {} - if self.DLDataReader.process_type == ProcessType.Simulation: - batch_grouped = batch.group_by( - ["obs_id", "event_id", "tel_type_id", "true_shower_primary_class"] - ) - elif self.DLDataReader.process_type == ProcessType.Observation: - batch_grouped = batch.group_by(["obs_id", "event_id", "tel_type_id"]) - features, mono_feature_vectors, stereo_feature_vectors = [], [], [] - true_shower_primary_class = [] - log_true_energy = [] - fov_lon, fov_lat, angular_separation = [], [], [] - cam_coord_offset_x, cam_coord_offset_y, cam_coord_distance = [], [], [] - for group_element in batch_grouped.groups: - if "features" in batch.colnames: - if self.sort_by_intensity: - # Sort images by the hillas intensity in a given batch if requested - group_element.sort(["hillas_intensity"], reverse=True) - # Stack the telescope images for stereo analysis - if self.stack_telescope_images: - # Retrieve the telescope images - plain_features = group_element["features"].data - # Stack the telescope images along the last axis - stacked_features = np.concatenate( - [plain_features[i] for i in range(plain_features.shape[0])], - axis=-1, - ) - # Append the stacked images to the features list - # shape: (batch_size, image_shape, image_shape, n_channels * n_tel) - features.append(stacked_features) - else: - # Append the plain images to the features list - # shape: (batch_size, n_tel, image_shape, image_shape, n_channels) - features.append(group_element["features"].data) - # Retrieve the feature vectors - if "mono_feature_vectors" in batch.colnames: - mono_feature_vectors.append(group_element["mono_feature_vectors"].data) - if "stereo_feature_vectors" in batch.colnames: - stereo_feature_vectors.append( - group_element["stereo_feature_vectors"].data - ) - # Retrieve the labels for the tasks - # FIXME: This won't work for divergent pointing directions - if "type" in self.tasks: - true_shower_primary_class.append( - group_element["true_shower_primary_class"].data[0] - ) - if "energy" in self.tasks: - log_true_energy.append(group_element["log_true_energy"].data[0]) - if "skydirection" in self.tasks: - fov_lon.append(group_element["fov_lon"].data[0]) - fov_lat.append( - group_element["fov_lat"].data[0] - ) - if "cameradirection" in self.tasks: - cam_coord_offset_x.append(group_element["cam_coord_offset_x"].data) - cam_coord_offset_y.append( - group_element["cam_coord_offset_y"].data - ) - # Store the labels in the labels dictionary - if "type" in self.tasks: - labels["type"] = to_categorical( - np.array(true_shower_primary_class), - num_classes=2, - ) - # Temp fix till keras support class weights for multiple outputs or I wrote custom loss - # https://github.com/keras-team/keras/issues/11735 - if len(self.tasks) == 1: - labels = to_categorical( - np.array(true_shower_primary_class), - num_classes=2, - ) - if "energy" in self.tasks: - labels["energy"] = np.array(log_true_energy) - if "skydirection" in self.tasks: - labels["skydirection"] = np.stack( - ( - np.array(fov_lon), - np.array(fov_lat), - ), - axis=1, - ) - if "cameradirection" in self.tasks: - labels["cameradirection"] = np.stack( - ( - np.array(cam_coord_offset_x), - np.array(cam_coord_offset_y), - ), - axis=1, - ) - # Store the fatures in the features dictionary - if "features" in batch.colnames: - features = np.array(features) - # TDOO: Add support for both feature vectors - if "mono_feature_vectors" in batch.colnames: - features = np.array(mono_feature_vectors) - if "stereo_feature_vectors" in batch.colnames: - features = np.array(stereo_feature_vectors) - return features, labels diff --git a/ctlearn/core/model.py b/ctlearn/core/model.py index bfac48c2..334b300d 100644 --- a/ctlearn/core/model.py +++ b/ctlearn/core/model.py @@ -7,15 +7,9 @@ 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", @@ -23,48 +17,6 @@ ] -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 - - class CTLearnModel(Component): """ Base component for creating a Keras model to be used in CTLearn. @@ -85,8 +37,8 @@ class CTLearnModel(Component): default_value={ "type": [512, 256, 2], "energy": [512, 256, 1], - "cameradirection": [512, 256, 2], - "skydirection": [512, 256, 2], + "cameradirection": [512, 256, 3], + "skydirection": [512, 256, 3], }, allow_none=False, help=( @@ -154,31 +106,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 +186,6 @@ class SingleCNN(CTLearnModel): def __init__( self, - input_shape, tasks, config=None, parent=None, @@ -251,114 +202,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 +264,6 @@ class ResNet(CTLearnModel): def __init__( self, - input_shape, tasks, config=None, parent=None, @@ -436,367 +283,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 +322,6 @@ class LoadedModel(CTLearnModel): def __init__( self, - input_shape, tasks, config=None, parent=None, @@ -842,54 +332,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..e69de29b diff --git a/ctlearn/core/pytorch/attention.py b/ctlearn/core/pytorch/attention.py new file mode 100644 index 00000000..a9408536 --- /dev/null +++ b/ctlearn/core/pytorch/attention.py @@ -0,0 +1,68 @@ +""" +This module defines the squeeze-excite blocks for channel-wise and/or spatial-wise attention mechanisms in PyTorch. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +__all__ = [ + "DualSqueezeExciteBlock", + "ChannelSqueezeExciteBlock", + "SpatialSqueezeExciteBlock", +] + +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/model.py b/ctlearn/core/pytorch/model.py new file mode 100644 index 00000000..a996240f --- /dev/null +++ b/ctlearn/core/pytorch/model.py @@ -0,0 +1,468 @@ +""" +This module defines the ``CTLearnModel`` classes, which holds the basic functionality for creating a PyTorch model to be used in CTLearn. +""" + +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, +) + +__all__ = [ + "BasicBlock", + "BottleneckBlock", + "MultiHeadClassifier", + "build_fully_connect_pytorch_head", + "PyTorchSingleCNN", + "PyTorchResNet", + "PyTorchLoadedModel", +] + + +class MultiHeadClassifier(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__() + # 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 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 the backbone output if it's spatially aggregated (B, C, 1, 1) -> (B, C) + if x.dim() > 2: + x = torch.flatten(x, start_dim=1) + + classification = None + energy = None + direction = None + + for original_task, internal_key in self._task_mapping.items(): + head = self.heads[internal_key] + out = head(x) + + if original_task == "type": + classification = F.softmax(out, dim=-1) + elif original_task == "energy": + energy = out + elif original_task in ["direction", "skydirection", "cameradirection"]: + direction = out + + return classification, energy, direction + + +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 MultiHeadClassifier(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) + + +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) + + # Global Average Pooling wrapper block + class GlobalAvgPool(nn.Module): + def forward(self, x): + return F.adaptive_avg_pool2d(x, (1, 1)) + + modules.append(GlobalAvgPool()) + + backbone_model = nn.Sequential(*modules) + return backbone_model, 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) + + # Global Average Pooling setup + class GlobalAvgPool(nn.Module): + def forward(self, x): + x = F.adaptive_avg_pool2d(x, (1, 1)) + return x.view(x.size(0), -1) + + modules.append(GlobalAvgPool()) + + 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`` handles loading a pre-saved PyTorch model weight layout. + """ + + def __init__(self, input_shape, tasks, config=None, parent=None, **kwargs): + super().__init__(tasks=tasks, config=config, parent=parent, **kwargs) + + # In PyTorch, instead of load_model returning an arbitrary configuration blindly, + # you instantiate the structure first and pass a state_dict or weights file path. + self.model = torch.load(self.load_model_from) + + if self.overwrite_head: + # Freeze/unfreeze backbone based on config choice + for param in self.model.backbone.parameters(): + param.requires_grad = self.trainable_backbone + + # Fetch out_features dynamically from existing backbone configuration + # This is a representative placeholder pattern for your architecture + out_features = self.head_layers[tasks[0]][0] + + self.logits_head = build_fully_connect_pytorch_head( + out_features, self.head_layers, self.head_activation_function, tasks + ) + self.model.head = self.logits_head \ No newline at end of file diff --git a/ctlearn/core/pytorch/model_collection.py b/ctlearn/core/pytorch/model_collection.py new file mode 100644 index 00000000..1cb1a41f --- /dev/null +++ b/ctlearn/core/pytorch/model_collection.py @@ -0,0 +1,206 @@ +""" +CTLearn PyTorch Model Registry + +This module defines the ctapipe Component wrapper classes for PyTorch models, +allowing them to be registered, configured, and instantiated dynamically using +ctapipe's Component system, matching the Keras model design. +""" + +from ctapipe.core import Component +from ctapipe.core.traits import Unicode, List, Float, Bool, Int, Dict, CaselessStrEnum +from ctlearn.core.ctlearn_enum import Task +import torch + +class CTLearnPyTorchModel(Component): + """ + Base class for PyTorch models in CTLearn. + Acts as a ctapipe Component wrapper for torch.nn.Module models. + """ + model_name = Unicode(help="Name of the model architecture").tag(config=True) + + def __init__(self, parent=None, **kwargs): + super().__init__(parent=parent, **kwargs) + self.model = None + + +class ThinResNet(CTLearnPyTorchModel): + """ + Component wrapper for ThinResNet PyTorch model. + """ + num_blocks = List( + trait=Int(), + default_value=[3, 4, 6, 3], + help="Number of blocks per stage in ResNet", + ).tag(config=True) + + dropout = Float( + default_value=0.1, + help="Dropout probability", + ).tag(config=True) + + use_bn = Bool( + default_value=False, + help="Whether to use Batch Normalization", + ).tag(config=True) + + def __init__(self, task="type", num_inputs=1, num_outputs=2, parent=None, **kwargs): + super().__init__(parent=parent, **kwargs) + from ctlearn.core.pytorch.nets.models.ThinResNet.ThinResNet import ThinResNet as PTThinResNet + self.model = PTThinResNet( + task=task, + num_inputs=num_inputs, + num_outputs=num_outputs, + num_blocks=self.num_blocks, + dropout=self.dropout, + use_bn=self.use_bn, + ) + + +class DoubleBBEfficientNet(CTLearnPyTorchModel): + """ + Component wrapper for DoubleBBEfficientNet PyTorch model. + """ + model_variant = Unicode( + default_value="efficientnet-b3", + help="Variant of EfficientNet backbone (e.g. efficientnet-b0 to b7)", + ).tag(config=True) + + def __init__(self, task="type", num_inputs=2, num_outputs=2, parent=None, **kwargs): + super().__init__(parent=parent, **kwargs) + from ctlearn.core.pytorch.nets.models.DoubleBBEfficientNet.DoubleBBEfficientNet import DoubleBBEfficientNet as PTDoubleBBEfficientNet + device_str = "cuda" + if parent is not None and hasattr(parent, "device_str"): + device_str = parent.device_str + self.model = PTDoubleBBEfficientNet( + task=task, + num_inputs=num_inputs, + num_outputs=num_outputs, + model_variant=self.model_variant, + device_str=device_str, + ) + + +class ThinResNet_DBB(CTLearnPyTorchModel): + """ + Component wrapper for ThinResNet_DBB (Dual Backbone ThinResNet) PyTorch model. + """ + num_blocks = List( + trait=Int(), + default_value=[3, 4, 6, 3], + help="Number of blocks per stage in ResNet", + ).tag(config=True) + + dropout = Float( + default_value=0.1, + help="Dropout probability", + ).tag(config=True) + + use_bn = Bool( + default_value=False, + help="Whether to use Batch Normalization", + ).tag(config=True) + + def __init__(self, task="type", num_inputs=2, num_outputs=3, parent=None, **kwargs): + super().__init__(parent=parent, **kwargs) + from ctlearn.core.pytorch.nets.models.ThinResNet_DBB.ThinResNet_DBB import ThinResNet_DBB as PTThinResNet_DBB + self.model = PTThinResNet_DBB( + task=task, + num_inputs=num_inputs, + num_outputs=num_outputs, + num_blocks=self.num_blocks, + dropout=self.dropout, + use_bn=self.use_bn, + ) + + +class PyTorchResNet(CTLearnPyTorchModel): + """ + Component wrapper for PyTorchResNet model (from custom_resnet_pytorch). + """ + init_layer = Dict( + default_value=None, + allow_none=True, + help="Parameters for the first convolutional layer.", + ).tag(config=True) + + init_max_pool = Dict( + default_value=None, + allow_none=True, + help="Parameters for the first max pooling layer.", + ).tag(config=True) + + residual_block_type = CaselessStrEnum( + ["basic", "bottleneck"], default_value="bottleneck", allow_none=False + ).tag(config=True) + + architecture = List( + trait=Dict(), + default_value=[ + {"filters": 48, "blocks": 2}, + {"filters": 96, "blocks": 3}, + {"filters": 128, "blocks": 3}, + {"filters": 256, "blocks": 3}, + ], + allow_none=False, + ).tag(config=True) + + init_padding = Int( + default_value=0, + allow_none=False, + min=0, + help="Initial padding to apply to the input data.", + ).tag(config=True) + + head_layers = Dict( + default_value={ + "type": [512, 256, 2], + "energy": [512, 256, 1], + "direction": [512, 256, 2], + }, + allow_none=False, + ).tag(config=True) + + head_activation_function = Dict( + default_value={ + "type": "relu", + "energy": "relu", + "direction": "tanh", + }, + allow_none=False, + ).tag(config=True) + + attention_mechanism = CaselessStrEnum( + ["Dual-SE", "Channel-SE", "Spatial-SE"], + default_value="Dual-SE", + allow_none=True, + ).tag(config=True) + + attention_reduction_ratio = Int( + default_value=16, + allow_none=True, + min=1, + ).tag(config=True) + + def __init__(self, task="type", num_inputs=1, num_outputs=2, parent=None, **kwargs): + super().__init__(parent=parent, **kwargs) + from ctlearn.core.pytorch.nets.models.PyTorchResNet.PyTorchResNet import PyTorchResNetModel + + # Override output layers dynamically if necessary based on num_outputs argument + head_layers = self.head_layers.copy() + if task in head_layers: + head_layers[task][-1] = num_outputs + + self.model = PyTorchResNetModel( + task=task, + num_inputs=num_inputs, + num_outputs=num_outputs, + init_padding=self.init_padding, + init_layer=self.init_layer, + init_max_pool=self.init_max_pool, + residual_block_type=self.residual_block_type, + architecture=self.architecture, + head_layers=head_layers, + head_activation_function=self.head_activation_function, + attention_mechanism=self.attention_mechanism, + attention_reduction_ratio=self.attention_reduction_ratio, + ) diff --git a/ctlearn/core/pytorch/net_utils.py b/ctlearn/core/pytorch/net_utils.py new file mode 100644 index 00000000..ff0ce0f4 --- /dev/null +++ b/ctlearn/core/pytorch/net_utils.py @@ -0,0 +1,608 @@ +""" +PyTorch Neural Network Utilities Module + +This module provides utility functions and helper classes for creating, managing, +and manipulating PyTorch neural network models in CTLearn. It includes functionality +for model creation, checkpoint management, visualization, and model export. + +Functions: + create_model: Factory function for instantiating models from configuration + +Classes: + ModelHelper: Collection of static utility methods for model operations +""" + +import importlib +import torch +import numpy as np +import os.path +import pickle +from matplotlib import pyplot as plt +import warnings +from ctlearn.core.ctlearn_enum import Task, Mode + +#------------------------------------------------------------------------------------------------------------------- +def create_model(model_parameters): + """ + Factory function to dynamically create and instantiate a model. + + This function uses Python's importlib to dynamically load a model class + based on the configuration parameters and instantiate it with the provided + parameters. This allows for flexible model selection without hardcoded imports. + + Args: + model_parameters (dict): Dictionary containing model configuration + Required keys: + - 'model_name' (str): Name of the model class to instantiate + Example: 'ResNet', 'EfficientNet', 'ViT' + - 'parameters' (dict): Dictionary of parameters to pass to model constructor + Example: {'num_outputs': 2, 'input_channels': 1, 'depth': 50} + + Returns: + torch.nn.Module: Instantiated model ready for training or inference + + Raises: + ValueError: If the model class is not found in the models module + ValueError: If there's an error instantiating the model (wrong parameters) + RuntimeError: If an unexpected error occurs during model creation + + Example: + >>> config = { + ... 'model_name': 'ResNet', + ... 'parameters': { + ... 'num_outputs': 2, + ... 'input_channels': 1, + ... 'depth': 50 + ... } + ... } + >>> model = create_model(config) + >>> print(type(model)) + + + Notes: + - All models must be located in ctlearn.core.pytorch.nets.models + - Model class name must match the module name + - Models should inherit from torch.nn.Module + """ + try: + # Define the base module path for models + module_name = "ctlearn.core.pytorch.nets.models" + model_type = model_parameters["model_name"] + model_params = model_parameters["parameters"] + + # Construct full class path (module.ModelClass) + full_class_path = f"ctlearn.core.pytorch.nets.models.{model_type}" + + # Dynamically import the model module + module = importlib.import_module(full_class_path) + # Get the module (intermediate step) + module = getattr(module, model_type) + # Get the actual model class from the module + model_class = getattr(module, model_type) + + # Instantiate the model with the provided parameters + model_net = model_class(**model_params) + return model_net + + except AttributeError: + raise ValueError(f"Model class {model_type} not found in module {module_name}.") + except TypeError as e: + raise ValueError(f"Error instantiating model {model_type}: {str(e)}") + except Exception as e: + raise RuntimeError(f"An unexpected error occurred: {str(e)}") +#------------------------------------------------------------------------------------------------------------------- + +class ModelHelper: + """ + Collection of static utility methods for PyTorch model operations. + + This class provides a suite of utility functions for common model operations + including parameter counting, serialization, visualization, kernel generation, + and model export. All methods are static and can be called without instantiation. + + Methods: + GetNumParamters: Count trainable parameters in a model + savePickle: Serialize data to pickle file + loadPickle: Deserialize data from pickle file + plotImage: Visualize tensor as image + GaborKernels: Generate Gabor filter kernels + saveModel: Save model weights to checkpoint + loadModel: Load model weights from checkpoint + exportOnnx: Export model to ONNX format + """ + + # ------------------------------------------------------------------------------------------------------------- + def GetNumParamters(self): + """ + Count the number of trainable parameters in the model. + + This method computes the total number of trainable parameters by summing + the number of elements in each parameter tensor that has requires_grad=True. + Useful for model complexity analysis and debugging. + + Returns: + tuple: (total_params, param_list) where: + - total_params (int): Total number of trainable parameters + - param_list (list): List of parameter counts per tensor + + Example: + >>> model = ResNet(num_outputs=2) + >>> helper = ModelHelper() + >>> total, per_layer = helper.GetNumParamters() + >>> print(f"Total trainable parameters: {total:,}") + Total trainable parameters: 23,512,130 + + Notes: + - Only counts parameters with requires_grad=True + - Frozen layers are not counted + - Includes biases if present + """ + numel_list = [ + p.numel() for p in self.model.parameters() if p.requires_grad == True + ] + return sum(numel_list), numel_list + + # ------------------------------------------------------------------------------------------------------------- + def savePickle(path, fileName, data): + """ + Save data to a pickle file. + + Serializes Python objects to a binary pickle file for later retrieval. + Useful for saving training metrics, configurations, or intermediate results. + + Args: + path (str): Directory path where the file will be saved + fileName (str): Name of the pickle file (should include .pkl extension) + data: Python object to serialize (can be dict, list, array, etc.) + + Example: + >>> metrics = {'loss': [0.5, 0.3, 0.2], 'accuracy': [0.7, 0.8, 0.85]} + >>> ModelHelper.savePickle('/path/to/results', 'metrics.pkl', metrics) + + Notes: + - File is opened in append binary mode ('ab') + - Multiple calls will append to the same file + - Use loadPickle to retrieve the data + """ + saveFile = os.path.join(path, fileName) + File = open(saveFile, "ab") + pickle.dump(data, File) + + # ------------------------------------------------------------------------------------------------------------- + def loadPickle(path, fileName): + """ + Load data from a pickle file. + + Deserializes a Python object that was previously saved with savePickle. + + Args: + path (str): Directory path where the file is located + fileName (str): Name of the pickle file to load + + Returns: + object: The deserialized Python object + + Raises: + FileNotFoundError: If the pickle file doesn't exist + pickle.UnpicklingError: If the file is corrupted or invalid + + Example: + >>> metrics = ModelHelper.loadPickle('/path/to/results', 'metrics.pkl') + >>> print(metrics['accuracy']) + [0.7, 0.8, 0.85] + """ + loadFile = os.path.join(path, fileName) + File = open(loadFile, "rb") + data = pickle.load(File) + return data + + # ------------------------------------------------------------------------------------------------------------- + def plotImage(img, permute=True): + """ + Display a tensor as an image using matplotlib. + + This utility function visualizes PyTorch tensors as grayscale images, + automatically handling gradient tracking and dimension permutation. + + Args: + img (torch.Tensor): Image tensor to display + Expected shapes: + - (H, W): Grayscale image + - (C, H, W): Multi-channel image (will be permuted to H, W, C) + permute (bool, optional): Whether to permute dimensions from (C, H, W) + to (H, W, C) for display. Defaults to True + + Side Effects: + - Displays image in matplotlib window + - Detaches tensor from computation graph if needed + + Example: + >>> # Display a model's first layer weights + >>> weights = model.conv1.weight[0, 0] # Single filter + >>> ModelHelper.plotImage(weights, permute=False) + + >>> # Display a preprocessed image + >>> img_tensor = batch['image'][0] # Shape: (1, 120, 120) + >>> ModelHelper.plotImage(img_tensor) + + Notes: + - Uses 'gray' colormap for visualization + - Automatically calls detach() for tensors in computation graph + - Blocking call - window must be closed to continue execution + """ + # Detach from computation graph if tensor is not a leaf + if img.is_leaf == False: + img = img.detach() + + # Permute from (C, H, W) to (H, W, C) for matplotlib + if permute and len(img.shape) == 3: + img = img.permute(1, 2, 0) + + plt.imshow(img, cmap="gray") + plt.show() + + # ------------------------------------------------------------------------------------------------------------- + def GaborKernels(size=7, showPlots=False): + """ + Generate a bank of Gabor filter kernels. + + Creates a set of Gabor filters with varying orientations and frequencies. + Gabor filters are useful for edge detection and texture analysis in images, + particularly for Cherenkov telescope images which have oriented features. + + Args: + size (int, optional): Size of the kernel (size x size). Defaults to 7 + showPlots (bool, optional): Whether to display each kernel. Defaults to False + + Returns: + list: List of numpy arrays, each representing a Gabor kernel + Length: 4 (orientations) × 5 (frequencies) = 20 kernels + Each kernel shape: (size, size) + + Kernel Parameters: + - Orientations (theta): 0°, 45°, 90°, 135° (4 angles) + - Frequencies: 0.15, 0.25, 0.35, 0.45, 0.55 (5 frequencies) + - Sigma: Fixed at 3 (spatial extent of the kernel) + + Example: + >>> # Generate 20 Gabor kernels of size 7x7 + >>> kernels = ModelHelper.GaborKernels(size=7, showPlots=False) + >>> print(f"Generated {len(kernels)} kernels") + Generated 20 kernels + >>> print(kernels[0].shape) + (7, 7) + + >>> # Visualize kernels during generation + >>> kernels = ModelHelper.GaborKernels(size=11, showPlots=True) + + Notes: + - Uses skimage.filters.gabor_kernel for generation + - Only real part of Gabor kernel is used + - Kernels are resized to specified size using bilinear interpolation + - Useful for initializing convolutional layers with oriented filters + + Applications: + - Initializing first convolutional layer weights + - Feature extraction for shower image analysis + - Edge and orientation detection in Cherenkov images + """ + try: + from skimage.filters import gabor_kernel + from skimage.transform import resize + except ImportError: + raise ImportError("skimage is required for GaborKernels. Install it via 'pip install scikit-image'") + + # prepare filter bank kernels + kernels = [] + # Iterate over 4 orientations + for theta in (0, np.pi / 4, np.pi / 2, 3 * np.pi / 4): + sigma = 3 # Fixed spatial extent + # Iterate over 5 frequencies + for frequency in (0.15, 0.25, 0.35, 0.45, 0.55): + # Generate Gabor kernel (complex-valued) + kernel = np.real( + gabor_kernel(frequency, theta=theta, sigma_x=sigma, sigma_y=sigma) + ) + + # Resize to specified size + kernel = resize(kernel, [size, size]) + kernels.append(kernel) + + # Optionally display each kernel + if showPlots: + print( + "Theta: ", + theta, + " Sigma: ", + sigma, + " Frequency: ", + frequency, + " Kernel size:", + kernel.shape, + ) + plt.imshow(kernel) + plt.show() + + return kernels + + # ------------------------------------------------------------------------------------------------------------- + def saveModel(model, data_path, filename): + """ + Save model weights to a checkpoint file. + + Saves the model's state dictionary (weights and biases) to a file + for later loading and inference or continued training. + + Args: + model (torch.nn.Module): The model to save + data_path (str): Directory path where the checkpoint will be saved + filename (str): Name of the checkpoint file (typically .pth extension) + + Side Effects: + - Creates a .pth file in the specified directory + - Prints confirmation message + + Example: + >>> model = ResNet(num_outputs=2) + >>> ModelHelper.saveModel(model, './checkpoints', 'best_model.pth') + Saving model: best_model.pth + + Notes: + - Only saves state_dict (weights), not the full model + - Does not save optimizer state or training history + - Use torch.save with full checkpoint dict for complete saving + """ + print("Saving model: ", filename) + torch.save(model.state_dict(), os.path.join(data_path, filename)) + + # ------------------------------------------------------------------------------------------------------------- + def loadModel(model, data_path, filename, mode, device_str='cpu'): + """ + Load model weights from a checkpoint file with robust error handling. + + This method loads pre-trained weights into a model, handling various + checkpoint formats and partial weight loading. It includes automatic + key name adjustment for different checkpoint structures and validates + weight dimensions. + + Args: + model (torch.nn.Module): The model to load weights into + data_path (str): Directory path where the checkpoint is located + filename (str): Name of the checkpoint file + mode (Mode): Operation mode (train, results, validate, observation, tunning) + Determines error handling strictness + device_str (str, optional): Device to load model onto ('cpu', 'cuda', 'cuda:0'). + Defaults to 'cpu' + + Returns: + torch.nn.Module: The model with loaded weights + + Checkpoint Format Compatibility: + Handles multiple checkpoint formats: + - Direct state_dict: {'layer.weight': tensor, ...} + - Wrapped state_dict: {'state_dict': {...}} + - Model state_dict: {'model_state_dict': {...}} + - Prefixed keys: {'model.0.layer.weight': tensor, ...} + + Key Matching Strategy: + 1. Remove 'model.0.' prefix from checkpoint keys if present + 2. Filter out keys not in model's state_dict + 3. Filter out keys with dimension mismatches + 4. Load only matching keys (strict=False) + + Error Handling: + - Training/Tunning mode: Issues warning for mismatches, continues + - Other modes: Raises ValueError for mismatches + - Missing checkpoint: Exits in non-training modes + + Example: + >>> model = ResNet(num_outputs=2) + >>> model = ModelHelper.loadModel( + ... model, + ... './checkpoints', + ... 'best_model.pth', + ... Mode.results, + ... device_str='cuda:0' + ... ) + Loading model: best_model.pth + Model Loaded. + + Notes: + - Uses weights_only=False for pickle compatibility (security warning) + - Automatically moves model to specified device + - Supports partial weight loading for transfer learning + - Strict=False allows loading subset of weights + + Security Warning: + Currently uses weights_only=False which can execute arbitrary code + during unpickling. Future versions should use weights_only=True. + """ + # Check if checkpoint file exists + if os.path.isfile(os.path.join(data_path, filename)): + print("Loading model: ", filename) + + # TODO: Test weights_only=True. Currently getting FutureWarning + # about security implications of weights_only=False + pretrained_dict = torch.load( + os.path.join(data_path, filename), + map_location=torch.device(device_str), + weights_only=False + ) + + # Handle different checkpoint formats + # Format 1: {'state_dict': {...}} + if type(pretrained_dict) == dict and "state_dict" in pretrained_dict: + pretrained_dict = pretrained_dict["state_dict"] + + # Format 2: {'model_state_dict': {...}} + if type(pretrained_dict) == dict and "model_state_dict" in pretrained_dict: + pretrained_dict = pretrained_dict["model_state_dict"] + + # Get current model's state dict + model_dict = model.state_dict() + + # Remove 'model.0.' prefix if present in checkpoint keys + modified_dict = {} + prefix = "model.0." + for key in pretrained_dict: + # Check if the key starts with 'model.0.' + if key.startswith(prefix): + # Remove the prefix and save with new key + new_key = key.replace(prefix, "") + modified_dict[new_key] = pretrained_dict[key] + + # Use modified dict if any keys were modified + if len(modified_dict) > 0: + pretrained_dict = modified_dict + + # Filter checkpoint to only include matching keys with same dimensions + pretrained_dict = { + k: v for k, v in pretrained_dict.items() + if k in model_dict and model_dict[k].size() == v.size() + } + + # Check for mismatches between checkpoint and model + if (len(model_dict) != len(pretrained_dict) or + set(model_dict.keys()) != set(pretrained_dict.keys())): + + pretrain_len = len(pretrained_dict) + model_len = len(model_dict) + # Keys only in checkpoint + unique_pretrained = set(pretrained_dict.keys()) - set(model_dict.keys()) + # Keys only in model + unique_model = set(model_dict.keys()) - set(pretrained_dict.keys()) + + # Strict error checking for non-training modes + if (mode != Mode.train and mode != Mode.tunning): + raise ValueError( + f"Error Loading the model. Pretrained Dict length: {pretrain_len} " + f"Model Dict length: {model_len}. Differences -> " + f"Pretrained keys: {unique_pretrained}, Model keys: {unique_model}" + ) + else: + # Warning for training/tuning modes (allow partial loading) + warnings.warn( + f"Warning Loading the model. Pretrained Dict length: {pretrain_len} " + f"Model Dict length: {model_len}. " + f"Differences -> Pretrained keys: {unique_pretrained}, " + f"Model keys: {unique_model}", + UserWarning + ) + + # Update model dict with pretrained weights + model_dict.update(pretrained_dict) + # Load the new state dict (strict=False allows partial loading) + model.load_state_dict(model_dict, strict=False) + + # Move model to specified device + device = torch.device(device_str) + model.to(device) + + print("Model Loaded.") + else: + # Checkpoint doesn't exist + model.to(torch.device(device_str)) + + print(f"CheckPoint file does not exist: {filename}") + # Exit if not in training mode (checkpoint required) + if mode != Mode.train: + exit() + + return model + + # ------------------------------------------------------------------------------------------------------------- + def exportOnnx(model, dummy_input, onnx_name, input_names, output_names): + """ + Export PyTorch model to ONNX format with simplification. + + Converts a PyTorch model to ONNX (Open Neural Network Exchange) format + for deployment and interoperability with other frameworks. Also applies + optimization and simplification to the exported model. + + Args: + model (torch.nn.Module): The PyTorch model to export + dummy_input (torch.Tensor or tuple): Example input for tracing + Must have same shape and type as model's expected input + onnx_name (str): Base name for output files (without extension) + input_names (list): Names for model inputs + Example: ['image', 'peak_time'] + output_names (list): Names for model outputs + Example: ['classification', 'energy', 'direction'] + + Side Effects: + - Creates two ONNX files: + 1. {onnx_name}.onnx - Original exported model + 2. {onnx_name}_simp.onnx - Simplified and optimized model + - Prints verbose export information + + Example: + >>> model = ResNet(num_outputs=2) + >>> model.eval() + >>> dummy_input = torch.randn(1, 1, 120, 120) + >>> ModelHelper.exportOnnx( + ... model, + ... dummy_input, + ... 'resnet_model', + ... input_names=['image'], + ... output_names=['classification'] + ... ) + + Requirements: + - pip install onnx + - pip install onnxsim + + Notes: + - Model must be in eval() mode before export + - Dummy input shape must match model's expected input + - Simplified model is validated before saving + - ONNX format allows deployment to: + * TensorRT (NVIDIA) + * OpenVINO (Intel) + * CoreML (Apple) + * ONNX Runtime (cross-platform) + + Simplification Benefits: + - Removes redundant operations + - Folds constant computations + - Optimizes graph structure + - Reduces model size + - Improves inference speed + """ + # Define dynamic axes to support variable batch sizes + dynamic_axes = {} + if input_names: + for name in input_names: + dynamic_axes[name] = {0: "batch_size"} + if output_names: + for name in output_names: + dynamic_axes[name] = {0: "batch_size"} + + # Export model to ONNX format + torch.onnx.export( + model, + dummy_input, + onnx_name + ".onnx", + verbose=True, + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + ) + + try: + import onnx + from onnxsim import simplify + except ImportError: + raise ImportError("onnx and onnxsim are required for exportOnnx. Install them via 'pip install onnx onnxsim'") + + # Load the exported ONNX model + model = onnx.load(onnx_name + ".onnx") + + # Apply simplification and optimization + model_simp, check = simplify(model) + + # Validate simplified model + assert check, "Simplified ONNX model could not be validated" + + # Save simplified model + onnx.save(model_simp, onnx_name + "_simp.onnx") + # ------------------------------------------------------------------------------------------------------------- diff --git a/ctlearn/core/pytorch/nets/__init__.py b/ctlearn/core/pytorch/nets/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/pytorch/nets/activation/__init__.py b/ctlearn/core/pytorch/nets/activation/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/pytorch/nets/activation/activation_functions.py b/ctlearn/core/pytorch/nets/activation/activation_functions.py new file mode 100644 index 00000000..a1fa3faf --- /dev/null +++ b/ctlearn/core/pytorch/nets/activation/activation_functions.py @@ -0,0 +1,70 @@ +""" +Neural Network Activation Functions Module + +This module provides custom activation functions and utilities for neural networks +in CTLearn. It includes both standard PyTorch activations and custom implementations +optimized for Cherenkov telescope data analysis. + +Activation functions introduce non-linearity into neural networks, allowing them +to learn complex patterns beyond simple linear transformations. The choice of +activation function can significantly impact model performance, training stability, +and convergence speed. + +Common Uses: + - ReLU and variants: Standard choice for hidden layers in CNNs + - Sigmoid: Binary classification output layers + - Tanh: Alternative to sigmoid with zero-centered output + - Softmax: Multi-class classification output layers + - Custom activations: Task-specific optimizations + +Imports: + torch: PyTorch tensor operations + torch.nn: Neural network modules and activation functions + +Example: + >>> import torch.nn as nn + >>> # Using standard PyTorch activations + >>> activation = nn.ReLU() + >>> x = torch.tensor([-1.0, 0.0, 1.0]) + >>> output = activation(x) # [0.0, 0.0, 1.0] + + >>> # In a neural network layer + >>> layer = nn.Sequential( + ... nn.Conv2d(1, 64, kernel_size=3), + ... nn.ReLU(), + ... nn.BatchNorm2d(64) + ... ) + +Notes: + - This module is a placeholder for future custom activation functions + - Standard PyTorch activations (nn.ReLU, nn.Sigmoid, etc.) are used throughout CTLearn + - Custom activations can be added here when needed for specific tasks + +Future Extensions: + - Swish/SiLU activation: x * sigmoid(x), shown to improve performance in some tasks + - GELU: Gaussian Error Linear Unit, used in transformers + - Mish: Self-regularized non-monotonic activation + - Parametric activations: PReLU, ELU with learnable parameters +""" + +import torch +import torch.nn as nn + +# This module currently serves as a placeholder for custom activation functions. +# Standard PyTorch activation functions are used directly from torch.nn: +# +# Common Activations Available: +# - nn.ReLU(): Rectified Linear Unit, f(x) = max(0, x) +# - nn.LeakyReLU(negative_slope): Leaky ReLU, f(x) = max(negative_slope*x, x) +# - nn.ELU(alpha): Exponential Linear Unit +# - nn.GELU(): Gaussian Error Linear Unit +# - nn.Sigmoid(): Sigmoid function, f(x) = 1 / (1 + exp(-x)) +# - nn.Tanh(): Hyperbolic tangent, f(x) = tanh(x) +# - nn.Softmax(dim): Softmax for multi-class classification +# - nn.LogSoftmax(dim): Log of softmax, numerically stable +# +# Usage Example: +# from torch.nn import ReLU, Sigmoid +# activation = ReLU() +# output = activation(input_tensor) + diff --git a/ctlearn/core/pytorch/nets/block/__init__.py b/ctlearn/core/pytorch/nets/block/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/pytorch/nets/block/cnn_blocks.py b/ctlearn/core/pytorch/nets/block/cnn_blocks.py new file mode 100644 index 00000000..a7d5bde1 --- /dev/null +++ b/ctlearn/core/pytorch/nets/block/cnn_blocks.py @@ -0,0 +1,344 @@ +""" +Convolutional Neural Network Building Blocks Module + +This module provides custom CNN building blocks for CTLearn neural network architectures. +It includes specialized layers for uncertainty quantification (evidential learning) and +residual blocks optimized for Cherenkov telescope image analysis. + +Classes: + Dirichlet: Evidential layer for classification with uncertainty quantification + NormalInvGamma: Evidential layer for regression with uncertainty quantification + ResBlock: Residual convolutional block with Gabor filter initialization +""" + +import torch +import torch.nn as nn +from ctlearn.core.pytorch.net_utils import ModelHelper +import torch.nn.functional as F + +class Dirichlet(nn.Module): + """ + Dirichlet distribution layer for evidential classification. + + This layer outputs parameters of a Dirichlet distribution, which is used + in evidential deep learning to quantify classification uncertainty. The + Dirichlet parameters (alphas) represent the concentration of probability + mass for each class, providing both predictions and uncertainty estimates. + + Mathematical Background: + Given evidence e_k for each class k, the Dirichlet parameters are: + α_k = e_k + 1 + + Where evidence is transformed through: + e_k = softplus(z_k) to ensure positivity + + Attributes: + dense (nn.Linear): Linear layer to compute raw evidence values + out_units (int): Number of output classes + + Methods: + evidence: Apply softplus activation to ensure positive evidence + forward: Compute Dirichlet parameters from input features + """ + + def __init__(self, in_features, out_units): + """ + Initialize the Dirichlet layer. + + Args: + in_features (int): Number of input features from previous layer + out_units (int): Number of output classes (Dirichlet dimensions) + """ + super().__init__() + self.dense = nn.Linear(in_features, out_units) + self.out_units = out_units + + def evidence(self, x): + """ + Transform raw outputs to positive evidence values. + + Uses softplus activation: log(1 + exp(x)) which is smooth and always positive. + This ensures that evidence values are strictly positive as required for + Dirichlet parameters. + + Args: + x (torch.Tensor): Raw evidence values from linear layer + + Returns: + torch.Tensor: Positive evidence values + """ + return F.softplus(x) + + def forward(self, x): + """ + Compute Dirichlet distribution parameters. + + This method transforms input features into Dirichlet parameters (alphas) + which characterize the predictive distribution over classes. Higher alpha + values indicate higher confidence in that class. + + Args: + x (torch.Tensor): Input features with shape (batch_size, in_features) + + Returns: + torch.Tensor: Dirichlet parameters (alphas) with shape (batch_size, out_units) + Each alpha_k > 1, where higher values indicate more evidence for class k + + Note: + The sum of alphas S = Σα_k represents total evidence. + Class probabilities can be computed as: p_k = α_k / S + Uncertainty can be quantified using various metrics on the Dirichlet distribution + """ + # Compute raw evidence + out = self.dense(x) + # Transform to Dirichlet parameters (add 1 to ensure α > 1) + alpha = self.evidence(out) + 1 + return alpha + +class NormalInvGamma(nn.Module): + """ + Normal Inverse Gamma distribution layer for evidential regression. + + This layer outputs parameters of a Normal Inverse Gamma (NIG) distribution, + used in evidential deep learning for regression with uncertainty quantification. + The NIG distribution provides both aleatoric (data) and epistemic (model) uncertainty. + + Mathematical Background: + The NIG distribution is parameterized by (μ, v, α, β): + - μ: Predicted mean + - v: Virtual observation count (epistemic uncertainty) + - α: Shape parameter (controls uncertainty distribution) + - β: Scale parameter (aleatoric uncertainty) + + Predictive variance: var = β / (v * (α - 1)) + + Attributes: + dense (nn.Linear): Linear layer to compute 4 parameters + out_units (int): Number of regression outputs (typically 1) + + Methods: + evidence: Apply softplus to ensure positive parameters + forward: Compute NIG parameters or predictions based on mode + """ + + def __init__(self, in_features, out_units): + """ + Initialize the Normal Inverse Gamma layer. + + Args: + in_features (int): Number of input features from previous layer + out_units (int): Number of regression outputs (usually 1) + """ + super().__init__() + # Output 4 parameters per regression target: μ, log(v), log(α), log(β) + self.dense = nn.Linear(in_features, out_units * 4) + self.out_units = out_units + + def evidence(self, x): + """ + Transform raw outputs to positive parameter values. + + Uses softplus activation to ensure v, α, and β are strictly positive + as required by the Normal Inverse Gamma distribution. + + Args: + x (torch.Tensor): Raw parameter values + + Returns: + torch.Tensor: Positive parameter values + """ + return F.softplus(x) + + def forward(self, x): + """ + Compute NIG distribution parameters or predictions. + + During training, returns all four NIG parameters for loss computation. + During inference, returns mean prediction and uncertainty estimate. + + Args: + x (torch.Tensor): Input features with shape (batch_size, in_features) + + Returns: + Training mode: + tuple: (mu, v, alpha, beta) - all NIG parameters + - mu: Mean predictions (batch_size, out_units) + - v: Virtual observations (batch_size, out_units) + - alpha: Shape parameters (batch_size, out_units) + - beta: Scale parameters (batch_size, out_units) + + Inference mode: + tuple: (mu, var) - predictions and uncertainties + - mu: Mean predictions (batch_size, out_units) + - var: Predictive variance (batch_size, out_units) + + Note: + The predictive variance combines both aleatoric and epistemic uncertainty: + var = β / (v * (α - 1)) + + Higher v means lower epistemic uncertainty (more confident) + Higher β means higher aleatoric uncertainty (noisier data) + """ + # Compute raw parameters + out = self.dense(x) + # Split into 4 components + mu, logv, logalpha, logbeta = torch.split(out, self.out_units, dim=-1) + + # Transform to ensure positivity + v = self.evidence(logv) + alpha = self.evidence(logalpha) + 1 # Ensure α > 1 + beta = self.evidence(logbeta) + + # Return appropriate outputs based on mode + if self.training: + # Return all parameters for evidential loss computation + return mu, v, alpha, beta + else: + # Return prediction and uncertainty for inference + # Predictive variance from NIG distribution + var = torch.sqrt(beta / (v * (alpha - 1))) + return mu, var + +class ResBlock(nn.Module): + """ + Residual convolutional block with Gabor filter initialization. + + This block implements a residual connection with convolutional layers, + batch normalization, and pooling. It's optimized for processing Cherenkov + telescope images by initializing filters with Gabor kernels that are + particularly effective at detecting oriented features in shower images. + + Architecture: + Input → Conv2d → BatchNorm → LeakyReLU → (+) → MaxPool → Conv1x1 → Dropout → Output + ↑ + | + Residual Connection + + Attributes: + conv (nn.Conv2d): Main convolutional layer + conv_dropout (nn.Dropout2d): Spatial dropout for regularization + batch_norm (nn.BatchNorm2d): Batch normalization layer + pool (nn.MaxPool2d): Max pooling layer (2x2) + activation (nn.LeakyReLU): Activation function + conv_out (nn.Conv2d): 1x1 conv for channel adjustment + + Methods: + forward: Process input through residual block + """ + + def __init__(self, n_chans_in, n_chans_out, kernel_size=3, conv_drop_pro=0.2): + """ + Initialize the residual block. + + Args: + n_chans_in (int): Number of input channels + n_chans_out (int): Number of output channels + kernel_size (int, optional): Size of convolutional kernel. Defaults to 3 + conv_drop_pro (float, optional): Dropout probability. Defaults to 0.2 + + Initialization Strategy: + 1. Kaiming normal initialization for main conv layer weights + 2. Batch norm weights initialized to 0.5 + 3. Batch norm biases initialized to 0 + 4. First few filters initialized with Gabor kernels for edge detection + """ + super(ResBlock, self).__init__() + + # Main convolutional layer (maintains channel dimension) + self.conv = nn.Conv2d( + n_chans_in, + n_chans_in, + kernel_size=kernel_size, + padding=int(kernel_size/2), # Same padding + bias=False # Bias not needed before batch norm + ) + + # Spatial dropout for regularization + self.conv_dropout = nn.Dropout2d(p=conv_drop_pro) + + # Batch normalization for training stability + self.batch_norm = nn.BatchNorm2d(num_features=n_chans_in) + + # Max pooling to reduce spatial dimensions + self.pool = nn.MaxPool2d(2) + + # LeakyReLU activation (allows small negative gradients) + self.activation = nn.LeakyReLU() + + # 1x1 convolution to adjust channel dimension + self.conv_out = nn.Conv2d( + n_chans_in, + n_chans_out, + kernel_size=1, + padding=0, + bias=False + ) + + # Initialize conv layer with Kaiming normal + # Appropriate for LeakyReLU activation + torch.nn.init.kaiming_normal_(self.conv.weight, nonlinearity='leaky_relu') + + # Initialize batch norm parameters + torch.nn.init.constant_(self.batch_norm.weight, 0.5) + torch.nn.init.zeros_(self.batch_norm.bias) + + # Initialize first filters with Gabor kernels + # Gabor filters are effective for detecting oriented features + # particularly useful for elongated shower images + kernels = ModelHelper.GaborKernels(size=kernel_size, showPlots=False) + for i in range(min(self.conv.weight.shape[0], len(kernels))): + with torch.no_grad(): + # Scale Gabor kernel by 100 for appropriate magnitude + self.conv.weight[i, :] = torch.nn.Parameter( + torch.tensor(kernels[i] * 100) + ) + + def forward(self, x): + """ + Forward pass through the residual block. + + Processing Steps: + 1. Apply convolution to extract features + 2. Normalize with batch norm + 3. Apply activation function + 4. Add residual connection (skip connection) + 5. Reduce spatial dimensions with pooling + 6. Adjust channels with 1x1 convolution + 7. Apply dropout for regularization + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, n_chans_in, height, width) + + Returns: + torch.Tensor: Output tensor with shape + (batch_size, n_chans_out, height/2, width/2) + + Note: + The residual connection helps gradient flow during backpropagation + and allows the network to learn identity mappings when beneficial. + + Spatial dimensions are halved due to max pooling with stride 2. + """ + # Convolutional feature extraction + out = self.conv(x) + + # Normalize activations + out = self.batch_norm(out) + + # Non-linear activation + out = self.activation(out) + + # Add residual connection (element-wise addition) + # This helps with gradient flow and allows learning identity mappings + out = out + x + + # Reduce spatial dimensions + out = self.pool(out) + + # Adjust number of channels + out = self.conv_out(out) + + # Apply dropout for regularization + out = self.conv_dropout(out) + + return out \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/loss_functions/__init__.py b/ctlearn/core/pytorch/nets/loss_functions/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/pytorch/nets/loss_functions/loss_functions.py b/ctlearn/core/pytorch/nets/loss_functions/loss_functions.py new file mode 100644 index 00000000..7f06c0e3 --- /dev/null +++ b/ctlearn/core/pytorch/nets/loss_functions/loss_functions.py @@ -0,0 +1,700 @@ +""" +Loss Functions Module for Deep Learning + +This module provides custom loss functions for various deep learning tasks in CTLearn, +with a focus on uncertainty quantification through evidential learning and specialized +losses for astronomical data analysis. + +Loss Functions: + evidential_regression_loss: Normal Inverse Gamma loss for regression with uncertainty + cosine_direction_loss: Cosine similarity loss for direction reconstruction + AngularDistance: Angular distance metric for celestial coordinates + AngularError: Angular error between vector predictions + VectorLoss: Angle-based loss for 2D vectors + FocalLoss: Focal loss for imbalanced classification + BCELogitsLoss: Binary cross-entropy with label smoothing + EvidClassification: Evidential classification with Dirichlet distribution + +Utility Functions: + smooth_BCE: Generate smoothed BCE targets + generate_hot_ones: Create one-hot encoded targets with smoothing +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class evidential_regression_loss(nn.Module): + """ + Evidential regression loss using Normal Inverse Gamma (NIG) distribution. + + This loss function implements evidential deep learning for regression tasks, + providing both point predictions and uncertainty estimates. It combines + a negative log-likelihood term with a regularization term to prevent + overconfident predictions. + + Mathematical Background: + The NIG distribution is parameterized by (μ, v, α, β): + - μ: Predicted mean + - v: Virtual observation count (inverse epistemic uncertainty) + - α: Shape parameter + - β: Scale parameter (related to aleatoric uncertainty) + + Total Loss = NLL + λ * Regularization + + Attributes: + lamb (float): Weight for the regularization term + reduction (str): Reduction method ('mean', 'sum', or None) + """ + + def __init__(self, lamb=1.0, reduction='mean'): + """ + Initialize the evidential regression loss. + + Args: + lamb (float, optional): Regularization weight. Defaults to 1.0 + Higher values encourage more conservative uncertainty estimates + reduction (str, optional): Reduction method. Defaults to 'mean' + Options: 'mean', 'sum', None + """ + super(evidential_regression_loss, self).__init__() + self.reduction = reduction + self.lamb = lamb + + def nig_nll(self, mu, v, alpha, beta, y): + """ + Compute the Negative Log-Likelihood for Normal Inverse Gamma distribution. + + This method calculates the NLL which measures how well the predicted + NIG distribution matches the observed target values. + + Args: + mu (torch.Tensor): Predicted mean values + v (torch.Tensor): Virtual observation counts + alpha (torch.Tensor): Shape parameters + beta (torch.Tensor): Scale parameters + y (torch.Tensor): Ground truth targets + + Returns: + torch.Tensor: Negative log-likelihood values + """ + two_beta_lambda = 2 * beta * (1 + v) + t1 = 0.5 * (torch.pi / v).log() + t2 = alpha * two_beta_lambda.log() + t3 = (alpha + 0.5) * (v * (y - mu) ** 2 + two_beta_lambda).log() + t4 = alpha.lgamma() + t5 = (alpha + 0.5).lgamma() + nll = t1 - t2 + t3 + t4 - t5 + return nll + + def nig_reg(self, mu, v, alpha, _beta, y): + """ + Compute the Normal Inverse Gamma regularization term. + + This regularization penalizes large prediction errors when the model + is highly confident (high evidence), encouraging the model to match + its uncertainty to its actual performance. + + Args: + mu (torch.Tensor): Predicted mean values + v (torch.Tensor): Virtual observation counts + alpha (torch.Tensor): Shape parameters + _beta (torch.Tensor): Scale parameters (not used) + y (torch.Tensor): Ground truth targets + + Returns: + torch.Tensor: Regularization values + + Raises: + RuntimeError: If reduction method is not 'sum', 'mean', or None + """ + # Regularization based on prediction error weighted by evidence + reg = (y - mu).abs() * (2 * v + alpha) + + if self.reduction == "mean": + error = reg.mean() + elif self.reduction == "sum": + error = reg.sum() + elif self.reduction == None or self.reduction == 'None': + error = reg + else: + raise RuntimeError("Reduction not supported: Use sum or mean") + + return error + + def set_lambda(self, lamb): + """ + Update the regularization weight. + + Args: + lamb (float): New regularization weight + """ + self.lamb = lamb + + def forward(self, dist_params, y): + """ + Compute the evidential regression loss. + + Args: + dist_params (tuple): Tuple of (mu, v, alpha, beta) parameters + y (torch.Tensor): Ground truth targets + + Returns: + torch.Tensor: Combined NLL and regularization loss + """ + # Unpack distribution parameters + if len(y) > 1: + mu, v, alpha, beta = (d.squeeze() for d in dist_params) + else: + mu, v, alpha, beta = (d for d in dist_params) + + # Compute regularization term + nig_reg_error = self.nig_reg(mu, v, alpha, beta, y) + + # Compute negative log-likelihood + nig_nll_error = self.nig_nll(mu, v, alpha, beta, y) + + # Apply reduction to NLL + if self.reduction == "mean": + nig_nll_error = nig_nll_error.mean() + elif self.reduction == "sum": + nig_nll_error = nig_nll_error.sum() + elif self.reduction == None or self.reduction == 'None': + nig_nll_error = nig_nll_error + else: + raise RuntimeError("Reduction not supported: Use sum or mean") + + # Combine NLL and weighted regularization + return nig_nll_error + self.lamb * nig_reg_error + +def cosine_direction_loss(pred_x, pred_y, true_x, true_y, reduction="mean"): + """ + Compute cosine similarity loss for 2D direction vectors. + + This loss function measures the angular difference between predicted and + true direction vectors using cosine similarity. It's particularly useful + for shower direction reconstruction where we care about the angle rather + than the magnitude. + + Mathematical Formula: + loss = 1 - cos(θ) = 1 - (pred · true) / (|pred| |true|) + + Args: + pred_x (torch.Tensor): Predicted x-components + pred_y (torch.Tensor): Predicted y-components + true_x (torch.Tensor): True x-components + true_y (torch.Tensor): True y-components + reduction (str, optional): Reduction method. Defaults to "mean" + Options: 'mean', 'sum', 'none' + + Returns: + torch.Tensor: Cosine direction loss + + Raises: + RuntimeError: If reduction method is not supported + + Note: + Both predicted and true vectors are normalized before computing + the dot product, making the loss independent of vector magnitude. + """ + # Normalize prediction vectors to unit length + pred_vec = F.normalize(torch.stack([pred_x, pred_y], dim=1), dim=1) + # Normalize true vectors to unit length + true_vec = F.normalize(torch.stack([true_x, true_y], dim=1), dim=1) + + # Compute 1 - cosine similarity (0 for perfect alignment, 2 for opposite directions) + if reduction == "mean": + return 1 - torch.sum(pred_vec * true_vec, dim=1).mean() + elif reduction == "sum": + return 1 - torch.sum(pred_vec * true_vec, dim=1).sum() + elif reduction == "none": + return 1 - torch.sum(pred_vec * true_vec, dim=1) + else: + raise RuntimeError("Reduction not supported: Use sum, mean or none") + +def AngularDistance(alt1_rad, alt2_rad, az1_rad, az2_rad, reduction=None): + """ + Calculate the angular distance between celestial coordinates. + + This function computes the great circle distance between two points on + the celestial sphere using the spherical law of cosines. Used for + evaluating direction reconstruction accuracy in astronomy. + + Mathematical Formula: + cos(Δθ) = cos(alt1)cos(alt2)cos(az1-az2) + sin(alt1)sin(alt2) + Δθ = arccos(cos(Δθ)) + + Args: + alt1_rad (torch.Tensor): Altitude of first points in radians + alt2_rad (torch.Tensor): Altitude of second points in radians + az1_rad (torch.Tensor): Azimuth of first points in radians + az2_rad (torch.Tensor): Azimuth of second points in radians + reduction (str, optional): Reduction method. Defaults to None + Options: 'sum', 'mean', None + + Returns: + tuple: (angular_distance_rad, angular_distance_deg) + - angular_distance_rad: Angular distances in radians + - angular_distance_deg: Angular distances in degrees + + Raises: + RuntimeError: If reduction method is not supported + + Note: + - Handles numerical edge cases (cosdelta = ±1) explicitly + - Clamps cosdelta to prevent arccos domain errors + - Returns both radians and degrees for convenience + """ + # Compute cosine of angular distance using spherical law of cosines + cosdelta = torch.cos(alt1_rad) * torch.cos(alt2_rad) * torch.cos(az1_rad - az2_rad) + \ + torch.sin(alt1_rad) * torch.sin(alt2_rad) + + # Clamp to valid range for arccos with small epsilon to avoid edge cases + cosdelta = torch.clamp(cosdelta, -1.0 + 1e-7, 1.0 - 1e-7) + + # Calculate angular distance in radians + ang_dist_rad = torch.acos(cosdelta) + + # Handle exact edge cases explicitly + ang_dist_rad[cosdelta == 1.0] = 0.0 # Perfect alignment + ang_dist_rad[cosdelta == -1.0] = torch.pi # Opposite directions + + # Convert to degrees + ang_dist_deg = torch.rad2deg(ang_dist_rad) + + # Apply reduction + if reduction == "sum": + return ang_dist_rad.sum(), ang_dist_deg.sum() + elif reduction == "mean": + return ang_dist_rad.mean(), ang_dist_deg.mean() + elif reduction == None or reduction == 'None': + return ang_dist_rad, ang_dist_deg + else: + raise RuntimeError("Reduction not supported: Use sum, mean or None") + +def AngularError(vec1, vec2, reduction='mean'): + """ + Compute angular error between 3D direction vectors. + + This function calculates the angle between two vectors using the dot product + formula. Useful for evaluating 3D direction predictions in Cartesian coordinates. + + Mathematical Formula: + cos(θ) = (v1 · v2) / (|v1| |v2|) + θ = arccos(cos(θ)) + + Args: + vec1 (torch.Tensor): First set of vectors with shape (batch_size, 3) + vec2 (torch.Tensor): Second set of vectors with shape (batch_size, 3) + reduction (str, optional): Reduction method. Defaults to 'mean' + Options: 'sum', 'mean', None + + Returns: + tuple: (angle_rad, angle_deg) + - angle_rad: Angular errors in radians + - angle_deg: Angular errors in degrees + + Raises: + RuntimeError: If reduction method is not supported + + Note: + - Handles vectors of any magnitude (not necessarily unit vectors) + - Clamps cosine values to prevent arccos domain errors + - Returns both radians and degrees + """ + # Compute dot product for each pair of vectors + dot_product = torch.sum(vec1 * vec2, dim=1) + + # Compute vector magnitudes (L2 norms) + norm_vec1 = torch.norm(vec1, dim=1) + norm_vec2 = torch.norm(vec2, dim=1) + + # Compute cosine of angle between vectors + cos_theta = dot_product / (norm_vec1 * norm_vec2) + + # Clamp to valid range for arccos + cos_theta = torch.clamp(cos_theta, -1.0, 1.0) + + # Compute angle in radians + angle_rad = torch.acos(cos_theta) + + # Convert to degrees + angle_deg = torch.rad2deg(angle_rad) + + # Apply reduction + if reduction == "sum": + return angle_rad.sum(), angle_deg.sum() + elif reduction == "mean": + return angle_rad.mean(), angle_deg.mean() + elif reduction == None or reduction == 'None': + return angle_rad, angle_deg + else: + raise RuntimeError("Reduction not supported: Use sum, mean or None") + + +class VectorLoss(nn.Module): + """ + Angle-based loss for 2D vector predictions. + + This loss function penalizes the angular difference between predicted and + target 2D vectors, regardless of their magnitudes. Useful when direction + matters more than magnitude. + + Attributes: + alpha (float): Scaling factor for the loss + reduction (str): Reduction method ('mean' or 'sum') + """ + + def __init__(self, alpha=0.001, reduction='mean'): + """ + Initialize the VectorLoss. + + Args: + alpha (float, optional): Scaling factor. Defaults to 0.001 + reduction (str, optional): Reduction method. Defaults to 'mean' + """ + super(VectorLoss, self).__init__() + self.alpha = alpha + self.reduction = reduction + + def forward(self, output, target): + """ + Compute the vector angle loss. + + This method calculates the absolute angular difference between + output and target vectors, normalized to [0, π]. + + Args: + output (torch.Tensor): Predicted vectors with shape (batch_size, 2) + target (torch.Tensor): Target vectors with shape (batch_size, 2) + + Returns: + torch.Tensor: Angular difference loss + + Raises: + RuntimeError: If reduction method is not 'sum' or 'mean' + """ + # Calculate angles using atan2 (returns range [-π, π]) + angles_output = torch.atan2(output[:, 1], output[:, 0]) + angles_target = torch.atan2(target[:, 1], target[:, 0]) + + # Compute absolute angle difference + angle_diff = torch.abs(angles_output - angles_target) + + # Normalize to [0, π] range (shortest angular path) + angle_diff = torch.remainder(angle_diff + torch.pi, 2 * torch.pi) - torch.pi + angle_diff = torch.abs(angle_diff) + + # Apply reduction + if self.reduction == "sum": + return angle_diff.sum() + elif self.reduction == "mean": + return angle_diff.mean() + else: + raise RuntimeError("Reduction not supported: Use sum or mean") + +def smooth_BCE(eps=0.1): + """ + Generate smoothed BCE (Binary Cross-Entropy) targets for label smoothing. + + Label smoothing prevents the model from becoming overconfident by using + soft targets instead of hard 0/1 labels. This can improve generalization. + + Args: + eps (float, optional): Smoothing parameter. Defaults to 0.1 + Determines how much to smooth the labels + + Returns: + tuple: (positive_label, negative_label) + - positive_label: Smoothed value for positive class (1 - 0.5*eps) + - negative_label: Smoothed value for negative class (0.5*eps) + + Example: + >>> cp, cn = smooth_BCE(eps=0.1) + >>> cp # 0.95 instead of 1.0 + >>> cn # 0.05 instead of 0.0 + """ + return 1.0 - 0.5 * eps, 0.5 * eps + + +def generate_hot_ones(device, cn, cp, outputs, targets): + """ + Create one-hot encoded targets with label smoothing. + + This function generates soft one-hot encoded targets for multi-class + classification with label smoothing applied. + + Args: + device: PyTorch device (CPU or CUDA) + cn (float): Negative class smoothed value + cp (float): Positive class smoothed value + outputs (torch.Tensor): Model outputs (used for shape) + targets (torch.Tensor): Ground truth class indices + + Returns: + torch.Tensor: Smoothed one-hot encoded targets + Shape: same as outputs + + Example: + >>> outputs = torch.randn(32, 10) # batch_size=32, num_classes=10 + >>> targets = torch.tensor([3, 7, 1, ...]) # class indices + >>> cp, cn = smooth_BCE(eps=0.1) + >>> t = generate_hot_ones(device, cn, cp, outputs, targets) + >>> # t[0] = [0.05, 0.05, 0.05, 0.95, 0.05, ...] # class 3 is positive + """ + # Initialize all values to negative class value + t = torch.full_like(outputs, cn, device=device) + n = outputs.shape[0] + # Set positive class values + t[range(n), targets] = cp + return t + +class FocalLoss(nn.Module): + """ + Focal Loss for addressing class imbalance in classification. + + Focal loss down-weights easy examples and focuses training on hard negatives. + This is particularly useful for highly imbalanced datasets where the model + can achieve high accuracy by simply predicting the majority class. + + Mathematical Formula: + FL(p_t) = -α_t (1-p_t)^γ log(p_t) + + where: + - p_t: model's estimated probability for the correct class + - α_t: weighting factor (alpha parameter) + - γ: focusing parameter (gamma parameter) + + Attributes: + alpha (torch.Tensor or None): Class weights + gamma (float): Focusing parameter (higher = more focus on hard examples) + reduction (str): Reduction method ('mean', 'sum', or 'none') + """ + + def __init__(self, alpha=None, gamma=2.0, reduction='mean'): + """ + Initialize the Focal Loss. + + Args: + alpha (torch.Tensor or None, optional): Class weights. Defaults to None + If None, all classes weighted equally + If Tensor, should have shape (num_classes,) + gamma (float, optional): Focusing parameter. Defaults to 2.0 + 0 = equivalent to cross-entropy + Higher values = more focus on hard examples + reduction (str, optional): Reduction method. Defaults to 'mean' + """ + super(FocalLoss, self).__init__() + self.alpha = alpha + self.gamma = gamma + self.reduction = reduction + + def set_alpha(self, alpha): + """ + Update the class weights. + + Args: + alpha (torch.Tensor): New class weights + """ + self.alpha = alpha + + def forward(self, inputs, targets): + """ + Compute the focal loss. + + Args: + inputs (torch.Tensor): Raw model outputs (logits) + Shape: (batch_size, num_classes) + targets (torch.Tensor): Ground truth class indices + Shape: (batch_size,) + + Returns: + torch.Tensor: Focal loss value + + Process: + 1. Compute standard cross-entropy loss + 2. Compute p_t (probability of correct class) + 3. Apply focal term: (1 - p_t)^gamma + 4. Weight by cross-entropy + 5. Apply class weights (alpha) if provided + """ + # Compute cross-entropy loss (unreduced) + ce_loss = F.cross_entropy(inputs, targets, weight=self.alpha, reduction='none') + + # Get probability of correct class + pt = torch.exp(-ce_loss) + + # Apply focal term and weight by CE loss + focal_loss = (1 - pt) ** self.gamma * ce_loss + + # Apply reduction + if self.reduction == 'mean': + return focal_loss.mean() + elif self.reduction == 'sum': + return focal_loss.sum() + else: + return focal_loss + + +class BCELogitsLoss(nn.Module): + """ + Binary Cross-Entropy with Logits and Label Smoothing. + + This loss combines binary cross-entropy with label smoothing for improved + generalization. It operates on raw logits (pre-sigmoid outputs). + + Attributes: + device: PyTorch device + label_smoothing (float): Label smoothing parameter + BCE: Binary cross-entropy loss function + cp (float): Smoothed positive label value + cn (float): Smoothed negative label value + """ + + def __init__(self, device, cls_pw=1.0, label_smoothing=0.0): + """ + Initialize the BCE with logits loss. + + Args: + device: PyTorch device (CPU or CUDA) + cls_pw (float, optional): Positive class weight. Defaults to 1.0 + label_smoothing (float, optional): Label smoothing factor. Defaults to 0.0 + """ + super().__init__() + self.device = device + self.label_smoothing = label_smoothing + self.BCE = nn.BCEWithLogitsLoss( + pos_weight=torch.tensor(cls_pw, device=self.device) + ) + # Generate smoothed label values + self.cp, self.cn = smooth_BCE(eps=self.label_smoothing) + + def forward(self, outputs, targets): + """ + Compute BCE loss with label smoothing. + + Args: + outputs (torch.Tensor): Raw model outputs (logits) + targets (torch.Tensor): Ground truth class indices + + Returns: + torch.Tensor: BCE loss value + """ + # Generate smoothed one-hot targets + t = generate_hot_ones(self.device, self.cn, self.cp, outputs, targets) + + # Compute BCE loss + bce_loss = self.BCE(outputs, t) + + return bce_loss + + +class EvidClassification(): + """ + Evidential classification using Dirichlet distribution. + + This class implements evidential deep learning for classification, + where the model outputs a Dirichlet distribution over class probabilities + rather than point estimates. This provides both predictions and uncertainty. + + Attributes: + class_weights (torch.Tensor or None): Optional class weights + """ + + def __init__(self, class_weights=None): + """ + Initialize evidential classification. + + Args: + class_weights (torch.Tensor or None, optional): Class weights. + Defaults to None + """ + self.class_weights = class_weights + + def dirichlet_reg(self, alpha, y): + """ + Compute Dirichlet distribution regularization. + + This regularization term encourages the model to output a uniform + Dirichlet distribution for incorrect classes, preventing overconfident + wrong predictions. + + Args: + alpha (torch.Tensor): Dirichlet parameters (concentration) + y (torch.Tensor): One-hot encoded targets + + Returns: + torch.Tensor: KL divergence from uniform Dirichlet + """ + # Remove evidence from correct class, keep evidence from wrong classes + alpha = y + (1 - y) * alpha + + # Uniform Dirichlet distribution (target) + beta = torch.ones_like(alpha) + + # Compute KL divergence between Dirichlet distributions + sum_alpha = alpha.sum(-1) + sum_beta = beta.sum(-1) + + t1 = sum_alpha.lgamma() - sum_beta.lgamma() + t2 = (alpha.lgamma() - beta.lgamma()).sum(-1) + t3 = alpha - beta + t4 = alpha.digamma() - sum_alpha.digamma().unsqueeze(-1) + + kl = t1 - t2 + (t3 * t4).sum(-1) + return kl.sum() + + def dirichlet_mse(self, alpha, y): + """ + Compute mean squared error for Dirichlet distribution. + + This term measures the accuracy of the mean prediction from the + Dirichlet distribution, accounting for its uncertainty. + + Args: + alpha (torch.Tensor): Dirichlet parameters + y (torch.Tensor): One-hot encoded targets + + Returns: + torch.Tensor: MSE loss + """ + # Sum of Dirichlet parameters (total evidence) + sum_alpha = alpha.sum(-1, keepdims=True) + + # Mean prediction (expected probability) + p = alpha / sum_alpha + + # Prediction error term + t1 = (y - p).pow(2) + + # Uncertainty term (variance of Dirichlet) + t2 = ((p * (1 - p)) / (sum_alpha + 1)) + + # Apply class weights if provided + if self.class_weights is not None: + t1 = t1 * self.class_weights.unsqueeze(0) + t2 = t2 * self.class_weights.unsqueeze(0) + + mse = t1 + t2 + return mse.sum() + + def loss(self, alpha, y, lamb=1.0): + """ + Compute the evidential classification loss. + + Combines Dirichlet MSE with regularization to produce predictions + with calibrated uncertainty estimates. + + Args: + alpha (torch.Tensor): Dirichlet parameters from model + y (torch.Tensor): Ground truth class indices + lamb (float, optional): Regularization weight. Defaults to 1.0 + + Returns: + torch.Tensor: Combined loss value + """ + num_classes = alpha.shape[-1] + # Convert to one-hot encoding + y = F.one_hot(y, num_classes) + # Combine MSE and regularization + return self.dirichlet_mse(alpha, y) + lamb * self.dirichlet_reg(alpha, y) diff --git a/ctlearn/core/pytorch/nets/models/DBBDanet/DBBDanet.py b/ctlearn/core/pytorch/nets/models/DBBDanet/DBBDanet.py new file mode 100644 index 00000000..ec67aad2 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/DBBDanet/DBBDanet.py @@ -0,0 +1,440 @@ +""" +Dual-Backbone DANet (Dual Attention Network) Module + +This module implements a dual-backbone architecture with attention mechanisms for +processing Cherenkov telescope images. It combines two separate backbones to process +image and timing information independently, then fuses the features for final prediction. + +The architecture incorporates both channel and spatial attention mechanisms to improve +feature representation and focus on the most relevant information in telescope images. + +Classes: + ChannelAttentionModule: Channel attention mechanism using global pooling + SpatialAttentionModule: Spatial attention mechanism using convolutional layers + DANet: Single backbone with dual attention mechanisms + DBBDanet: Dual-backbone network combining two DANet architectures + +References: + - DANet: "Dual Attention Network for Scene Segmentation" (CVPR 2019) + - CBAM: "Convolutional Block Attention Module" (ECCV 2018) +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class ChannelAttentionModule(nn.Module): + """ + Channel Attention Module using global pooling. + + This module learns to emphasize informative channels and suppress less useful ones + by exploiting inter-channel relationships. It uses both average and max pooling to + capture different aspects of channel-wise statistics. + + Architecture: + Input → [AvgPool, MaxPool] → Shared MLP → Element-wise Sum → Sigmoid → Channel Weights + + Attributes: + avg_pool (nn.AdaptiveAvgPool2d): Global average pooling + max_pool (nn.AdaptiveMaxPool2d): Global max pooling + fc (nn.Sequential): Shared MLP for channel attention + sigmoid (nn.Sigmoid): Activation for attention weights + """ + + def __init__(self, in_channels, reduction=16): + """ + Initialize the Channel Attention Module. + + Args: + in_channels (int): Number of input channels + reduction (int, optional): Channel reduction ratio for the MLP bottleneck. + Defaults to 16. Higher values reduce parameters but may lose information. + + Example: + >>> cam = ChannelAttentionModule(in_channels=512, reduction=16) + >>> x = torch.randn(8, 512, 7, 7) + >>> out = cam(x) # Same shape as input, but with channel attention applied + """ + super(ChannelAttentionModule, self).__init__() + + # Global pooling operations to capture channel-wise statistics + self.avg_pool = nn.AdaptiveAvgPool2d(1) # Output: (B, C, 1, 1) + self.max_pool = nn.AdaptiveMaxPool2d(1) # Output: (B, C, 1, 1) + + # Shared MLP: Channel reduction → ReLU → Channel restoration + self.fc = nn.Sequential( + nn.Conv2d(in_channels, in_channels // reduction, 1, bias=False), + nn.ReLU(), + nn.Conv2d(in_channels // reduction, in_channels, 1, bias=False) + ) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + """ + Apply channel attention to the input feature map. + + Process: + 1. Apply global average pooling and max pooling separately + 2. Pass both through shared MLP + 3. Sum the two attention maps + 4. Apply sigmoid to get attention weights in [0, 1] + 5. Multiply with original input (element-wise) + + Args: + x (torch.Tensor): Input feature map with shape (B, C, H, W) + B: batch size, C: channels, H: height, W: width + + Returns: + torch.Tensor: Channel-attended feature map with same shape as input + Important channels are emphasized, less important ones suppressed + """ + # Process through average pooling path + avg_out = self.fc(self.avg_pool(x)) # (B, C, 1, 1) + + # Process through max pooling path + max_out = self.fc(self.max_pool(x)) # (B, C, 1, 1) + + # Combine both paths and apply sigmoid + out = avg_out + max_out # (B, C, 1, 1) + + # Apply attention weights to input + return self.sigmoid(out) * x # (B, C, H, W) + +class SpatialAttentionModule(nn.Module): + """ + Spatial Attention Module using channel pooling. + + This module learns to focus on important spatial locations in the feature map + by exploiting inter-spatial relationships. It uses both average and max pooling + across channels to generate a spatial attention map. + + Architecture: + Input → [AvgPool(channel), MaxPool(channel)] → Concat → Conv → Sigmoid → Spatial Weights + + Attributes: + conv (nn.Conv2d): Convolutional layer to generate spatial attention + sigmoid (nn.Sigmoid): Activation for attention weights + """ + + def __init__(self, kernel_size=7): + """ + Initialize the Spatial Attention Module. + + Args: + kernel_size (int, optional): Size of convolutional kernel. Defaults to 7. + Larger kernels capture wider spatial context but increase computation. + Common choices: 3, 5, 7 + + Example: + >>> sam = SpatialAttentionModule(kernel_size=7) + >>> x = torch.randn(8, 512, 7, 7) + >>> out = sam(x) # Same shape, but with spatial attention applied + """ + super(SpatialAttentionModule, self).__init__() + + # Calculate padding to maintain spatial dimensions + padding = kernel_size // 2 + + # Convolution to process concatenated spatial statistics + # Input: 2 channels (avg + max), Output: 1 channel (attention map) + self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + """ + Apply spatial attention to the input feature map. + + Process: + 1. Compute average across channels + 2. Compute max across channels + 3. Concatenate the two maps + 4. Apply convolution to generate spatial attention map + 5. Apply sigmoid to get attention weights in [0, 1] + 6. Multiply with original input (element-wise) + + Args: + x (torch.Tensor): Input feature map with shape (B, C, H, W) + + Returns: + torch.Tensor: Spatially-attended feature map with same shape as input + Important spatial locations are emphasized + """ + # Average pooling across channel dimension + avg_out = torch.mean(x, dim=1, keepdim=True) # (B, 1, H, W) + + # Max pooling across channel dimension + max_out, _ = torch.max(x, dim=1, keepdim=True) # (B, 1, H, W) + + # Concatenate spatial statistics + x_cat = torch.cat([avg_out, max_out], dim=1) # (B, 2, H, W) + + # Generate spatial attention map + x_att = self.conv(x_cat) # (B, 1, H, W) + + # Apply attention weights to input + return self.sigmoid(x_att) * x # (B, C, H, W) + +class DANet(nn.Module): + """ + Dual Attention Network (DANet) - Single backbone with attention mechanisms. + + This network implements a CNN backbone with both channel and spatial attention + mechanisms. It progressively reduces spatial dimensions while increasing channel + depth, then applies dual attention before classification. + + Architecture: + Input → Conv1 → BN → ReLU → MaxPool → + Layer1 (64→128) → Layer2 (128→256) → Layer3 (256→512) → + Channel Attention → Spatial Attention → + Global AvgPool → Dropout → FC + + Attributes: + conv1 (nn.Conv2d): Initial convolution layer + bn1 (nn.BatchNorm2d): Batch normalization + relu (nn.ReLU): Activation function + maxpool (nn.MaxPool2d): Max pooling layer + layer1, layer2, layer3: Feature extraction blocks + cam (ChannelAttentionModule): Channel attention + sam (SpatialAttentionModule): Spatial attention + avgpool (nn.AdaptiveAvgPool2d): Global average pooling + dropout (nn.Dropout): Dropout for regularization + fc (nn.Linear): Final classification layer + """ + + def __init__(self, num_inputs=1, num_classes=2, dropout_rate=0.3): + """ + Initialize the DANet model. + + Args: + num_inputs (int, optional): Number of input channels. Defaults to 1. + For telescope images: 1 for charge only, 2 for charge+timing + num_classes (int, optional): Number of output classes. Defaults to 2. + For particle classification: 2 (gamma vs proton) + For regression: 1 + dropout_rate (float, optional): Dropout probability. Defaults to 0.3. + Higher values = more regularization but may underfit + """ + super(DANet, self).__init__() + + # Initial convolution: large kernel for receptive field + self.conv1 = nn.Conv2d(num_inputs, 64, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + + # Progressive feature extraction layers + self.layer1 = self._make_layer(64, 128, 2) # 64 → 128 channels + self.layer2 = self._make_layer(128, 256, 2) # 128 → 256 channels + self.layer3 = self._make_layer(256, 512, 2) # 256 → 512 channels + + # Dual attention modules + self.cam = ChannelAttentionModule(512) # Channel attention on 512 channels + self.sam = SpatialAttentionModule() # Spatial attention + + # Final classification layers + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) # Global pooling to (B, 512, 1, 1) + self.dropout = nn.Dropout(p=dropout_rate) + self.fc = nn.Linear(512, num_classes) + + def _make_layer(self, in_channels, out_channels, blocks): + """ + Create a sequence of convolutional blocks. + + Each block consists of: Conv → BatchNorm → ReLU + The first block changes channel dimension, subsequent blocks maintain it. + + Args: + in_channels (int): Number of input channels + out_channels (int): Number of output channels + blocks (int): Number of convolutional blocks + + Returns: + nn.Sequential: Sequential container of conv blocks + """ + layers = [] + # First block: change channel dimension + layers.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)) + layers.append(nn.BatchNorm2d(out_channels)) + layers.append(nn.ReLU(inplace=True)) + + # Remaining blocks: maintain channel dimension + for _ in range(1, blocks): + layers.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)) + layers.append(nn.BatchNorm2d(out_channels)) + layers.append(nn.ReLU(inplace=True)) + + return nn.Sequential(*layers) + + def forward(self, x): + """ + Forward pass through the DANet. + + Args: + x (torch.Tensor): Input tensor with shape (B, C_in, H, W) + B: batch size, C_in: input channels (1 or 2) + H, W: image dimensions (typically 120x120) + + Returns: + torch.Tensor: Output predictions with shape (B, num_classes) + For classification: logits (pre-softmax) + For regression: predicted values + """ + # Initial convolution and downsampling + x = self.conv1(x) # (B, 64, H/2, W/2) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) # (B, 64, H/4, W/4) + + # Feature extraction layers + x = self.layer1(x) # (B, 128, H/4, W/4) + x = self.layer2(x) # (B, 256, H/4, W/4) + x = self.layer3(x) # (B, 512, H/4, W/4) + + # Apply dual attention mechanisms with residual connections + x = self.cam(x) + x # Channel attention + residual + x = self.sam(x) + x # Spatial attention + residual + + # Global pooling and classification + x = self.avgpool(x) # (B, 512, 1, 1) + x = torch.flatten(x, 1) # (B, 512) + x = self.dropout(x) # Regularization + x = self.fc(x) # (B, num_classes) + + return x + +class DBBDanet(nn.Module): + """ + Dual-Backbone DANet for multi-modal telescope data processing. + + This architecture uses two separate DANet backbones to process different + input modalities (e.g., charge and timing information) independently, + then fuses their features for final prediction. This allows each backbone + to specialize in extracting relevant features from its input modality. + + Fusion Strategies: + - Concatenation: Preserves all information but doubles feature dimension + - Addition: Reduces dimension but may lose information + + Attributes: + task (str): Task type ('type', 'energy', 'direction') + use_concat (bool): Whether to concatenate or add backbone outputs + backbone_1 (DANet): First backbone for primary input (charge) + backbone_2 (DANet): Second backbone for secondary input (timing) + fc (nn.Linear): Final classification/regression layer + dropout (nn.Dropout): Dropout for regularization + """ + + def __init__(self, task, num_inputs=1, num_classes=2, use_concat=False, dropout_rate=0.3): + """ + Initialize the Dual-Backbone DANet. + + Args: + task (str): Task type to perform + Options: 'type' (classification), 'energy' (regression), 'direction' (regression) + num_inputs (int, optional): Number of input channels per backbone. Defaults to 1. + Typically 1 (grayscale images) + num_classes (int, optional): Number of output classes/values. Defaults to 2. + For 'type': 2 (gamma, proton) + For 'energy': 1 (energy value) + For 'direction': 2 or 3 (angular coordinates) + use_concat (bool, optional): Whether to concatenate backbone outputs. + Defaults to False (use addition instead) + True: More parameters, preserves all information + False: Fewer parameters, may lose some information + dropout_rate (float, optional): Dropout probability. Defaults to 0.3. + + Example: + >>> # For particle classification with concatenation + >>> model = DBBDanet(task='type', num_inputs=1, num_classes=2, use_concat=True) + >>> + >>> # For energy regression with addition + >>> model = DBBDanet(task='energy', num_inputs=1, num_classes=1, use_concat=False) + """ + super(DBBDanet, self).__init__() + + self.task = task + self.use_concat = use_concat + + # Initialize two separate DANet backbones + self.backbone_1 = DANet(num_inputs=num_inputs, num_classes=num_classes, dropout_rate=dropout_rate) + self.backbone_2 = DANet(num_inputs=num_inputs, num_classes=num_classes, dropout_rate=dropout_rate) + + # Get feature dimension from backbone + num_features = self.backbone_1.fc.in_features # Typically 512 + + # Adjust feature dimension based on fusion strategy + if self.use_concat: + num_features *= 2 # Double if concatenating + + # Create new final layers + self.fc = nn.Linear(num_features, num_classes) + self.dropout = nn.Dropout(p=dropout_rate, inplace=True) + + # Remove original final layers from backbones (use as feature extractors) + self.backbone_1.fc = nn.Identity() + self.backbone_1.dropout = nn.Identity() + + self.backbone_2.fc = nn.Identity() + self.backbone_2.dropout = nn.Identity() + + def forward(self, x): + if x.shape[1] >= 2: + x, y = torch.split(x, [1, x.shape[1]-1], dim=1) + else: + y = x + """ + Forward pass through the dual-backbone network. + + Process: + 1. Extract features from both inputs using separate backbones + 2. Fuse features (concatenate or add) + 3. Apply dropout and final classification/regression layer + 4. Return task-specific output + + Args: + x (torch.Tensor): First input (charge image) with shape (B, C, H, W) + y (torch.Tensor): Second input (timing image) with shape (B, C, H, W) + Both inputs should have the same spatial dimensions + + Returns: + tuple: (classification, energy, direction) where: + - classification: Class logits if task=='type', else None + - energy: Energy prediction if task=='energy', else None + - direction: Direction prediction if task=='direction', else None + Only one of the three is non-None based on self.task + + Note: + The dual-backbone design allows the model to learn separate + representations for different input modalities before fusion, + which can be more effective than early fusion approaches. + """ + # Initialize outputs (only one will be non-None) + energy = None + classification = None + direction = None + + # Extract features from both backbones independently + feature_1 = self.backbone_1(x) # Features from charge image + feature_2 = self.backbone_2(y) # Features from timing image + + # Fuse features based on selected strategy + if self.use_concat: + # Concatenate: Preserves all information + out = torch.cat((feature_1, feature_2), dim=1) # (B, 2*num_features) + else: + # Addition: Element-wise fusion + out = feature_1 + feature_2 # (B, num_features) + + # Apply regularization and final layer + out = self.dropout(out) + out = self.fc(out) # (B, num_classes) + + # Assign output based on task type + if self.task == "type": + classification = out # Classification logits + elif self.task == "energy": + energy = out # Energy prediction + elif self.task == "direction": + direction = out # Direction prediction + + return classification, energy, direction \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/DBBNoPropDTReg.py b/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/DBBNoPropDTReg.py new file mode 100644 index 00000000..3d6aa68f --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/DBBNoPropDTReg.py @@ -0,0 +1,406 @@ +""" +Dual-Backbone NoProp-DT Regression Model Module + +This module implements a dual-backbone denoising diffusion model for regression tasks +on Cherenkov telescope data. It uses a progressive denoising approach with multiple +diffusion steps to refine predictions, particularly effective for energy and direction +reconstruction tasks. + +The NoProp-DT (No-Propagation Denoising Transformer) architecture uses: +- Progressive denoising through T diffusion steps +- Cosine noise schedule for stable training +- SNR (Signal-to-Noise Ratio) weighting for loss computation +- Dual-backbone design for processing charge and timing information + +Classes: + DBBNoPropDTReg: Dual-backbone diffusion model for regression tasks + +References: + - "Denoising Diffusion Probabilistic Models" (Ho et al., NeurIPS 2020) + - "Improved Denoising Diffusion Probabilistic Models" (Nichol & Dhariwal, 2021) +""" + +import torch +from torch import nn +from .denoiseBlockThinRestNet import DenoiseBlock, MemoryEfficientSwish +import math + +class DBBNoPropDTReg(nn.Module): + """ + Dual-Backbone NoProp-DT model for regression tasks with diffusion. + + This model implements a denoising diffusion approach for regression where + predictions are progressively refined through multiple diffusion steps. + Each step removes noise from the target embedding, ultimately producing + a clean prediction. + + Architecture Overview: + Input (x, y) → [DenoiseBlock_1 → ... → DenoiseBlock_T] → Regressor → Output + + At each step t: + 1. Add noise to target embedding + 2. Pass through denoise block + 3. Refine embedding + 4. Final step: regress to prediction + + Key Components: + - DenoiseBlocks: T sequential denoising steps + - Target Embedder: Projects targets to latent space + - Regressor: Final prediction head + - Noise Schedule: Cosine schedule for adding noise + - SNR Weighting: Signal-to-noise ratio for loss weighting + + Attributes: + task (str): Task type ('energy', 'direction') + num_classes (int): Number of output values + embedding_dim (int): Dimension of latent embedding space + T (int): Number of diffusion steps + eta (float): Learning rate scaling factor + blocks (nn.ModuleList): List of denoising blocks + regressor (nn.Sequential): Final regression head + classifier (nn.Linear): Classification head (unused in regression) + target_embedder (nn.Linear): Projects targets to embedding space + alpha_bar (torch.Tensor): Noise schedule parameters + snr_diff (torch.Tensor): SNR difference for loss weighting + """ + + def __init__(self, task, num_outputs, embedding_dim=128, T=3, eta=0.1): + """ + Initialize the DBBNoPropDTReg model. + + Args: + task (str): Task type to perform + Options: 'energy' (energy regression), 'direction' (direction regression) + num_outputs (int): Number of regression outputs + For energy: 1 (single value) + For direction: 2 or 3 (angular coordinates) + embedding_dim (int, optional): Dimension of latent embedding space. + Defaults to 128. Higher values allow more complex representations + but increase memory usage. + T (int, optional): Number of diffusion steps. Defaults to 3. + More steps allow finer refinement but increase computation. + Typical values: 3-10 + eta (float, optional): Learning rate scaling factor. Defaults to 0.1. + Controls the weight of denoising loss vs final regression loss. + + Note: + The model uses Xavier uniform initialization for regressor weights + to ensure stable training at initialization. + """ + super().__init__() + + self.task = task + num_classes = num_outputs + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.T = T + self.eta = eta + + # Create T denoising blocks for progressive refinement + self.blocks = nn.ModuleList([ + DenoiseBlock(embedding_dim, num_channels=1) + for _ in range(T) + ]) + + # Final regression head with intermediate activation + # Architecture: Linear → Swish → Linear + # This provides non-linearity while maintaining differentiability + self.regressor = nn.Sequential( + nn.Linear(embedding_dim, embedding_dim // 2), + MemoryEfficientSwish(), # Memory-efficient activation function + nn.Linear(embedding_dim // 2, num_outputs) + ) + + # Classification head (included for architecture compatibility but unused) + self.classifier = nn.Linear(embedding_dim, num_classes) + + # Noise schedule: determines how much noise to add at each step + # Uses cosine schedule for smooth noise progression + self.register_buffer('alpha_bar', self._cosine_schedule(T)) + + # SNR differences for loss weighting + # Weights each diffusion step by its contribution to final quality + self.register_buffer('snr_diff', self._calculate_snr_diff(self.alpha_bar)) + + # Target embedder: projects ground truth values to latent space + self.target_embedder = nn.Linear(num_outputs, embedding_dim) + + # Initialize regressor weights with Xavier uniform + # Ensures initial gradients are neither too large nor too small + for m in self.regressor: + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + nn.init.zeros_(m.bias) + + def _cosine_schedule(self, T): + """ + Generate a cosine noise schedule for diffusion. + + This schedule determines how much noise is added at each diffusion step. + The cosine schedule provides a smooth progression from high noise to low noise, + which has been shown to improve training stability and final performance. + + Mathematical Formula: + α̅_t = cos²((t/T + 0.008) / 1.008 × π/2) + + where: + - t ∈ [1, T]: current diffusion step + - α̅_t: proportion of original signal retained + - (1 - α̅_t): proportion of noise + + Args: + T (int): Total number of diffusion steps + + Returns: + torch.Tensor: Alpha bar values for each step with shape (T,) + Values range from ~1.0 (step 1, low noise) to ~0.0 (step T, high noise) + + Properties: + - Monotonically decreasing: α̅_1 > α̅_2 > ... > α̅_T + - Smooth transitions between steps + - Small offset (0.008) prevents numerical issues at boundaries + + Example: + >>> schedule = self._cosine_schedule(5) + >>> print(schedule) + tensor([0.9950, 0.9801, 0.9553, 0.9211, 0.8782]) + """ + # Create timestep array from 1 to T + t = torch.arange(1, T + 1, dtype=torch.float32) + + # Apply cosine schedule formula + # The small offsets (0.008, 1.008) prevent division by zero and ensure + # alpha_bar doesn't reach exactly 0 or 1 + alpha_bar = torch.cos((t / T + 0.008) / 1.008 * (math.pi / 2)) ** 2 + + return alpha_bar + + def _calculate_snr_diff(self, alpha_bar): + """ + Calculate Signal-to-Noise Ratio differences for loss weighting. + + The SNR difference quantifies how much the signal quality improves + from one diffusion step to the next. This is used to weight the + denoising loss at each step - steps that make larger improvements + receive higher weight. + + Mathematical Formula: + SNR_t = α̅_t / (1 - α̅_t) + SNR_diff_t = SNR_t - SNR_{t-1} + + where: + - SNR_t: Signal-to-noise ratio at step t + - α̅_t: Alpha bar value at step t + + Args: + alpha_bar (torch.Tensor): Alpha bar values from noise schedule + Shape: (T,) + + Returns: + torch.Tensor: SNR differences with shape (T,) + Positive values indicate signal improvement + Clamped to minimum of 1e-5 to prevent numerical issues + + Implementation Details: + - SNR_0 is set to 0 (no signal before first step) + - Differences are clamped to ensure positive weights + - Small epsilon (1e-8) added to denominator for numerical stability + + Example: + >>> alpha_bar = torch.tensor([0.99, 0.95, 0.90]) + >>> snr_diff = self._calculate_snr_diff(alpha_bar) + >>> print(snr_diff) + tensor([99.0000, 19.0000, 9.0000]) # Approximate values + """ + # Calculate SNR at each step: α̅_t / (1 - α̅_t) + # Small epsilon prevents division by zero when alpha_bar ≈ 1 + snr = alpha_bar / (1 - alpha_bar + 1e-8) + + # Prepend SNR_0 = 0 (before first denoising step) + # Then compute differences: SNR_t - SNR_{t-1} + snr_prev = torch.cat([torch.tensor([0.]), snr[:-1]]) + + # Clamp to minimum value to ensure positive weights + # This prevents negative or zero weights that would destabilize training + return torch.clamp(snr - snr_prev, min=1e-5) + + def forward_denoise(self, x, y, z_prev, t): + """ + Perform one step of denoising. + + This method applies the t-th denoising block to refine the current + embedding estimate. Each block processes the input features along + with the current noisy embedding to produce a cleaner estimate. + + Args: + x (torch.Tensor): Primary input features (charge images) + Shape: (batch_size, channels, height, width) + y (torch.Tensor): Secondary input features (timing images) + Shape: (batch_size, channels, height, width) + z_prev (torch.Tensor): Previous/current embedding estimate + Shape: (batch_size, embedding_dim) + Contains noise that will be reduced in this step + t (int): Current diffusion step index (0 to T-1) + + Returns: + torch.Tensor: Refined embedding after denoising + Shape: (batch_size, embedding_dim) + Has less noise than z_prev + + Process: + 1. Denoise block extracts features from x and y + 2. Current embedding z_prev is used as query + 3. Block outputs refined embedding with reduced noise + + Note: + The fourth parameter (None) is for optional label embeddings, + which are not used in regression tasks. + """ + # Apply t-th denoising block + # Returns tuple (denoised_embedding, attention_weights) + # We only need the embedding, so take index [0] + return self.blocks[t](x, y, z_prev, None)[0] + + def regress(self, z): + """ + Generate final prediction from clean embedding. + + This method applies the regression head to convert the final + clean embedding into the actual prediction values (energy or direction). + + Args: + z (torch.Tensor): Clean embedding from final denoising step + Shape: (batch_size, embedding_dim) + + Returns: + torch.Tensor: Final predictions + Shape: (batch_size, num_outputs) + For energy: (batch_size, 1) + For direction: (batch_size, 2) or (batch_size, 3) + + Architecture: + embedding_dim → embedding_dim/2 → num_outputs + with Swish activation in between + + Example: + >>> z = torch.randn(32, 128) # Batch of 32, embedding_dim=128 + >>> pred = model.regress(z) + >>> print(pred.shape) # torch.Size([32, 1]) for energy + """ + return self.regressor(z) + + def inference(self, x, y): + """ + Perform full inference through all diffusion steps. + + This method executes the complete denoising process, starting from + random noise (or zeros during evaluation) and progressively refining + the embedding through T diffusion steps, finally producing a prediction. + + Process: + 1. Initialize embedding z: + - Training: Random Gaussian noise + - Inference: Zeros (deterministic) + 2. For each diffusion step t = 0 to T-1: + - Apply denoising block to refine z + 3. Final step: Apply regressor to get prediction + + Args: + x (torch.Tensor): Primary input features (charge images) + Shape: (batch_size, 1, height, width) + y (torch.Tensor): Secondary input features (timing images) + Shape: (batch_size, 1, height, width) + + Returns: + torch.Tensor: Final predictions + Shape: (batch_size, num_outputs) + + Behavior Difference: + Training (self.training=True): + - Starts from random noise + - Introduces stochasticity for exploration + - Helps learn robust denoising + + Evaluation (self.training=False): + - Starts from zeros + - Deterministic predictions + - More stable and reproducible + + Example: + >>> model.eval() # Set to evaluation mode + >>> x = torch.randn(8, 1, 120, 120) # Batch of 8 images + >>> y = torch.randn(8, 1, 120, 120) + >>> predictions = model.inference(x, y) + >>> print(predictions.shape) # torch.Size([8, 1]) for energy + """ + # Get batch size from input + B = x.size(0) + + # Initialize embedding + # Training: Random noise for stochastic exploration + # Evaluation: Zeros for deterministic predictions + z = torch.randn(B, self.embedding_dim, device=x.device) + if not self.training: + z = torch.zeros(B, self.embedding_dim, device=x.device) + + # Progressive denoising through T steps + for t in range(self.T): + z = self.forward_denoise(x, y, z, t) + + # Generate final prediction from clean embedding + return self.regress(z) + + def forward(self, x): + if x.shape[1] >= 2: + x, y = torch.split(x, [1, x.shape[1]-1], dim=1) + else: + y = x + """ + Forward pass through the model. + + This method provides a task-specific interface to the model, returning + predictions in the expected format for each task type. It wraps the + inference method and formats outputs appropriately. + + Args: + x (torch.Tensor): Primary input features (charge images) + Shape: (batch_size, 1, height, width) + y (torch.Tensor): Secondary input features (timing images) + Shape: (batch_size, 1, height, width) + + Returns: + tuple: (classification, energy, direction) where: + - classification: None (not used for regression) + - energy: Predictions if task=='energy', else None + - direction: Predictions if task=='direction', else None + + Only one of energy/direction is non-None based on self.task + + Task Routing: + - task == 'direction': Returns (None, None, predictions) + where predictions = (batch_size, 2 or 3) for angular coordinates + + - task == 'energy': Returns (None, predictions, None) + where predictions = (batch_size, 1) for energy values + + Example: + >>> model = DBBNoPropDTReg(task='energy', num_outputs=1) + >>> x = torch.randn(16, 1, 120, 120) + >>> y = torch.randn(16, 1, 120, 120) + >>> cls, energy, direction = model(x, y) + >>> print(cls) # None + >>> print(energy.shape) # torch.Size([16, 1]) + >>> print(direction) # None + + Note: + The tuple format (classification, energy, direction) is maintained + for compatibility with multi-task training frameworks, even though + this model only performs one task at a time. + """ + # Route to appropriate output based on task + if self.task == "direction": + # Direction reconstruction: return predictions in third position + return None, None, self.inference(x, y) + elif self.task == "energy": + # Energy regression: return predictions in second position + return None, self.inference(x, y), None diff --git a/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/denoiseBlockThinRestNet.py b/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/denoiseBlockThinRestNet.py new file mode 100644 index 00000000..2491d3e2 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/denoiseBlockThinRestNet.py @@ -0,0 +1,132 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn1 = AdaptiveBatchNorm2d(out_channels) + self.act1 = nn.GELU() + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn2 = AdaptiveBatchNorm2d(out_channels) + self.act2 = nn.GELU() + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + def forward(self, x): + # residual = self.act1(self.bn1(self.conv1(x))) + # residual = self.act2(self.bn2(self.conv2(residual))) + residual = self.act1((self.conv1(x))) + residual = self.act2((self.conv2(residual))) + residual = self.se(residual) + out = residual + self.shortcut(x) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, drop_prob=0.2): + super().__init__() + # Ahora más profundo y ancho: + self.conv_path_x = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + ) + + self.conv_path_y = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + ) + + self.conv_path_end = nn.Sequential( + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(256, 512), # Embedding más grande + nn.BatchNorm1d(512), + nn.GELU(), + ) + + self.fc_z1 = nn.Linear(embedding_dim, 512) + self.bn_z1 = nn.BatchNorm1d(512) + self.fc_z2 = nn.Linear(512, 512) + self.bn_z2 = nn.BatchNorm1d(512) + self.fc_z3 = nn.Linear(512, 512) + self.bn_z3 = nn.BatchNorm1d(512) + self.fc_f1 = nn.Linear(1024, 512) + self.bn_f1 = nn.BatchNorm1d(512) + self.fc_f2 = nn.Linear(512, 256) + self.bn_f2 = nn.BatchNorm1d(256) + self.fc_out = nn.Linear(256, embedding_dim) + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, y, z_prev, _): + x_feat = self.conv_path_x(x) + y_feat = self.conv_path_y(y) + + x_feat = self.conv_path_end(x_feat + y_feat) + + + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/denoiseBlockThinRestNet_original.py b/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/denoiseBlockThinRestNet_original.py new file mode 100644 index 00000000..fc52d51f --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/DBBNoPropDTReg/denoiseBlockThinRestNet_original.py @@ -0,0 +1,154 @@ +# Denoising block +import torch +from torch import nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +# class ResidualBlock(nn.Module): +# def __init__(self, in_channels, out_channels): +# super().__init__() +# self.conv_block = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels), +# nn.PReLU(), +# nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels) +# ) + +# self.shortcut = nn.Sequential() +# if in_channels != out_channels: +# self.shortcut = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=1), +# nn.BatchNorm2d(out_channels) +# ) + +# self.relu = nn.PReLU() + +# def forward(self, x): +# return self.relu(self.conv_block(x) + self.shortcut(x)) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16, drop_path_rate=0.1): + super().__init__() + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.norm1 = nn.GroupNorm(8, out_channels) + self.act1 = nn.GELU() + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + self.norm2 = nn.GroupNorm(8, out_channels) + self.act2 = nn.GELU() + + self.se = SEBlock(out_channels, reduction) + + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1), + nn.GroupNorm(8, out_channels) + ) + + try: + from timm.models.layers import DropPath + self.drop_path = DropPath(drop_path_rate) + except ImportError: + self.drop_path = nn.Identity() + + # self.final_act = nn.GELU() + self.final_act = MemoryEfficientSwish() + def forward(self, x): + residual = self.act1(self.norm1(self.conv1(x))) + residual = self.act2(self.norm2(self.conv2(residual))) + # residual = self.act1((self.conv1(x))) + # residual = self.act2((self.conv2(residual))) + + residual = self.se(residual) + + out = self.drop_path(residual) + self.shortcut(x) + return self.final_act(out) + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1): + super().__init__() + + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 32), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(32, 64), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(64, 128), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(128, 256), + ) + + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + self.fc_f1 = nn.Linear(512, 256) + self.bn_f1 = nn.BatchNorm1d(256) + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + self.fc_out = nn.Linear(128, embedding_dim) + + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + + z_feat = h3 + h1 + + h_f = torch.cat([x_feat, z_feat], dim=1) + + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/DBBRegNet/DBBRegNet.py b/ctlearn/core/pytorch/nets/models/DBBRegNet/DBBRegNet.py new file mode 100644 index 00000000..0d454ba1 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/DBBRegNet/DBBRegNet.py @@ -0,0 +1,318 @@ +""" +Dual-Backbone RegNet Model Module + +This module implements a dual-backbone architecture using RegNet (Regularized Networks) +as feature extractors for processing Cherenkov telescope data. RegNet is a family of +efficient networks designed with design principles derived from network design spaces. + +The dual-backbone approach allows processing charge and timing information independently +before fusing features for final predictions. + +Classes: + SingleChannelRegNet: Single RegNet backbone for one input modality + DBBRegNet: Dual-backbone architecture combining two RegNet backbones + +References: + - "Designing Network Design Spaces" (Radosavovic et al., CVPR 2020) + - RegNet: https://arxiv.org/abs/2003.13678 +""" + +import torch.nn.functional as F +import torch +from torchvision import models +import torch.nn as nn + + +class SingleChannelRegNet(nn.Module): + """ + Single-channel RegNet backbone for feature extraction. + + This class wraps a pretrained RegNet model from torchvision and adapts it + for single-channel telescope images. The first convolutional layer is modified + to accept custom input channels, and the final layer is adjusted for the + desired number of outputs. + + RegNet Architecture: + - Stage-based design with regularized block patterns + - Efficient depth and width distributions + - Better accuracy-efficiency trade-offs than EfficientNet + + Available RegNet Variants: + - regnet_y_400mf: ~4M parameters, 400MFLOPs + - regnet_y_800mf: ~6M parameters, 800MFLOPs (default) + - regnet_y_1_6gf: ~11M parameters, 1.6GFLOPs + - regnet_x_16gf: ~54M parameters, 16GFLOPs + + Attributes: + regnet (models.RegNet): Modified RegNet model from torchvision + """ + + def __init__(self, num_inputs=1, num_classes=2): + """ + Initialize the single-channel RegNet backbone. + + Args: + num_inputs (int, optional): Number of input channels. Defaults to 1. + For telescope data: 1 (grayscale charge or timing image) + num_classes (int, optional): Number of output classes/values. Defaults to 2. + For classification: 2 (gamma vs proton) + For regression: 1 (energy or direction components) + + Modifications: + 1. First conv layer: Modified to accept num_inputs channels + Original: Conv2d(3, 32, ...) for RGB images + Modified: Conv2d(num_inputs, 32, ...) for grayscale + + 2. Final layer: Modified to output num_classes values + Original: fc layer for ImageNet (1000 classes) + Modified: fc layer for custom task + + Example: + >>> # Create backbone for grayscale images, binary classification + >>> backbone = SingleChannelRegNet(num_inputs=1, num_classes=2) + >>> x = torch.randn(8, 1, 120, 120) # Batch of 8 grayscale images + >>> out = backbone(x) # Shape: (8, 2) + """ + super(SingleChannelRegNet, self).__init__() + + # Load pretrained RegNet-Y-800MF + # RegNet-Y variants have squeeze-and-excitation (SE) blocks + self.regnet = models.regnet_y_800mf(weights=models.RegNet_Y_800MF_Weights.DEFAULT) + + # Alternative RegNet variants (commented out): + # regnet_y_400mf: Smaller, faster, less accurate + # self.regnet = models.regnet_y_400mf(weights=models.RegNet_Y_400MF_Weights.DEFAULT) + + # regnet_x_16gf: Much larger, slower, potentially more accurate + # self.regnet = models.regnet_x_16gf(weights=models.RegNet_X_16GF_Weights.DEFAULT) + + # regnet_x_1_6gf: Medium size without SE blocks + # self.regnet = models.regnet_x_1_6gf(weights=models.RegNet_X_1_6GF_Weights.DEFAULT) + + # regnet_y_1_6gf: Medium size with SE blocks + # self.regnet = models.regnet_y_1_6gf(weights=models.RegNet_Y_1_6GF_Weights.DEFAULT) + + # Modify first convolutional layer for custom input channels + # Original stem: Conv2d(3, 32) for RGB images + # Modified stem: Conv2d(num_inputs, 32) for grayscale/custom + self.regnet.stem[0] = nn.Conv2d( + num_inputs, 32, + kernel_size=(3, 3), + stride=(2, 2), + padding=(1, 1), + bias=False + ) + + # Modify final fully connected layer for custom number of outputs + num_features = self.regnet.fc.in_features # Get size from pretrained model + self.regnet.fc = nn.Linear(num_features, num_classes) + + def forward(self, x): + """ + Forward pass through RegNet. + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, num_inputs, H, W) + Typically (batch_size, 1, 120, 120) for telescope images + + Returns: + torch.Tensor: Output predictions with shape (batch_size, num_classes) + For classification: logits before softmax + For regression: predicted values + + Architecture Flow: + Input → Stem (Conv+BN+ReLU) → + Stage 1 (depth=1) → Stage 2 (depth=3) → + Stage 3 (depth=7) → Stage 4 (depth=12) → + AvgPool → FC → Output + """ + return self.regnet(x) + +class DBBRegNet(nn.Module): + """ + Dual-Backbone RegNet for multi-modal telescope data. + + This architecture uses two separate RegNet backbones to process different + input modalities (e.g., charge and timing images) independently, then fuses + their features for final prediction. This allows each backbone to specialize + in extracting relevant features from its input modality. + + Fusion Strategies: + - Concatenation (use_concat=True): Preserves all information + Output features: 2 × num_features + Pros: Retains distinct information from both streams + Cons: Doubles parameter count in final layer + + - Addition (use_concat=False): Reduces dimension + Output features: num_features + Pros: Fewer parameters, forces feature alignment + Cons: May lose complementary information + + Attributes: + use_concat (bool): Whether to concatenate or add backbone outputs + task (str): Task type ('type', 'energy', 'direction') + bb1 (SingleChannelRegNet): First backbone for primary input + bb2 (SingleChannelRegNet): Second backbone for secondary input + dropout (nn.Dropout): Dropout layer for regularization + fc (nn.Linear): Final classification/regression layer + """ + + def __init__(self, task, use_concat=False, num_inputs=1, num_classes=2, dropout_rate=0.1): + """ + Initialize the Dual-Backbone RegNet. + + Args: + task (str): Task type to perform + Options: 'type' (classification), 'energy' (regression), 'direction' (regression) + Case-insensitive, will be converted to lowercase + use_concat (bool, optional): Whether to concatenate backbone outputs. + Defaults to False (uses addition) + True: Concatenate features (more parameters, preserves information) + False: Add features (fewer parameters, forces alignment) + num_inputs (int, optional): Number of input channels per backbone. Defaults to 1. + num_classes (int, optional): Number of output classes/values. Defaults to 2. + For 'type': 2 (gamma, proton) + For 'energy': 1 (energy value) + For 'direction': 2 or 3 (angular coordinates) + dropout_rate (float, optional): Dropout probability. Defaults to 0.1. + Applied before final layer for regularization + + Architecture: + Input_1 → Backbone_1 ↘ + → Fusion → Dropout → FC → Output + Input_2 → Backbone_2 ↗ + + Example: + >>> # Particle classification with concatenation + >>> model = DBBRegNet(task='type', use_concat=True, num_classes=2) + >>> charge = torch.randn(16, 1, 120, 120) + >>> timing = torch.randn(16, 1, 120, 120) + >>> cls, energy, direction = model(charge, timing) + >>> print(cls.shape) # torch.Size([16, 2]) + + >>> # Energy regression with addition + >>> model = DBBRegNet(task='energy', use_concat=False, num_classes=1) + >>> cls, energy, direction = model(charge, timing) + >>> print(energy.shape) # torch.Size([16, 1]) + """ + super(DBBRegNet, self).__init__() + + # Store configuration + self.use_concat = use_concat + self.task = task.lower() # Normalize to lowercase for consistency + + # Initialize two independent RegNet backbones + self.bb1 = SingleChannelRegNet(num_inputs=num_inputs, num_classes=num_classes) + self.bb2 = SingleChannelRegNet(num_inputs=num_inputs, num_classes=num_classes) + + # Get feature dimension from backbone + num_features = self.bb1.regnet.fc.in_features + + # Adjust feature dimension based on fusion strategy + if self.use_concat: + num_features *= 2 # Double for concatenation + + # Remove original classification heads from backbones + # Use backbones as feature extractors only + self.bb1.regnet.fc = nn.Identity() # Replace fc with identity (no-op) + self.bb2.regnet.fc = nn.Identity() + + # Regularization layer + self.dropout = nn.Dropout(p=dropout_rate, inplace=True) + + # Final task-specific prediction layer + self.fc = nn.Linear(num_features, num_classes) + + def forward(self, x): + if x.shape[1] >= 2: + x, y = torch.split(x, [1, x.shape[1]-1], dim=1) + else: + y = x + """ + Forward pass through the dual-backbone network. + + Process: + 1. Extract features from both inputs using separate backbones + 2. Fuse features (concatenate or add) + 3. Apply dropout for regularization + 4. Apply final layer for task-specific prediction + 5. Route output to appropriate task variable + + Args: + x (torch.Tensor): First input (charge image) with shape (B, C, H, W) + B: batch size, C: channels (typically 1), H/W: image dimensions + y (torch.Tensor): Second input (timing image) with same shape as x + + Returns: + tuple: (classification, energy, direction) where: + - classification: Class logits if task=='type', else None + Shape: (batch_size, 2) for binary classification + - energy: Energy prediction if task=='energy', else None + Shape: (batch_size, 1) for energy regression + - direction: Direction prediction if task=='direction', else None + Shape: (batch_size, 2) or (batch_size, 3) for angular coordinates + + Only one of the three is non-None based on self.task + + Feature Extraction Details: + - Both backbones process their inputs independently + - No gradient flow between backbones during feature extraction + - Each backbone can learn modality-specific representations + + Fusion Details: + Concatenation mode (use_concat=True): + out = [features_1 | features_2] # Shape: (B, 2F) + Retains all information from both modalities + + Addition mode (use_concat=False): + out = features_1 + features_2 # Shape: (B, F) + Forces features to be in same space + Acts as implicit alignment/fusion + + Example: + >>> model = DBBRegNet(task='type', use_concat=True) + >>> x = torch.randn(32, 1, 120, 120) # Charge images + >>> y = torch.randn(32, 1, 120, 120) # Timing images + >>> cls, energy, direction = model(x, y) + >>> + >>> # Only classification is non-None + >>> assert cls is not None + >>> assert energy is None + >>> assert direction is None + >>> print(cls.shape) # torch.Size([32, 2]) + """ + # Initialize outputs (only one will be non-None) + energy = None + classification = None + direction = None + + # Extract features from both backbones independently + feature_1 = self.bb1(x) # Features from charge image + feature_2 = self.bb2(y) # Features from timing image + + # Fuse features based on selected strategy + if self.use_concat: + # Concatenate features along feature dimension + # Preserves all information from both modalities + out = torch.cat((feature_1, feature_2), dim=1) # Shape: (B, 2F) + else: + # Element-wise addition of features + # Forces features into shared representation space + out = feature_1 + feature_2 # Shape: (B, F) + + # Apply dropout for regularization + out = self.dropout(out) + + # Apply final prediction layer + out = self.fc(out) # Shape: (B, num_classes) + + # Route output to appropriate task-specific variable + if self.task == "type": + classification = out # Classification logits + elif self.task == "energy": + energy = out # Energy predictions + elif self.task == "direction": + direction = out # Direction predictions + + # Return tuple format (for compatibility with multi-task frameworks) + return classification, energy, direction \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/DoubleBBEfficientNet/DoubleBBEfficientNet.py b/ctlearn/core/pytorch/nets/models/DoubleBBEfficientNet/DoubleBBEfficientNet.py new file mode 100644 index 00000000..9ba9ff00 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/DoubleBBEfficientNet/DoubleBBEfficientNet.py @@ -0,0 +1,474 @@ +""" +Dual-Backbone EfficientNet Model Module + +This module implements a dual-backbone architecture using EfficientNet models +as feature extractors for processing Cherenkov telescope data. EfficientNet +is a family of efficient neural networks that achieve state-of-the-art accuracy +with fewer parameters through compound scaling. + +The dual-backbone approach processes charge and timing information independently +before fusing features for final predictions, with specialized heads for different tasks. + +Classes: + MemoryEfficientSwish: Memory-efficient Swish activation function + SEBlock: Squeeze-and-Excitation attention block + DoubleBBEfficientNet: Dual-backbone EfficientNet for multi-task learning + +References: + - EfficientNet: "EfficientNet: Rethinking Model Scaling for CNNs" (ICML 2019) + - Squeeze-and-Excitation: "Squeeze-and-Excitation Networks" (CVPR 2018) +""" + +import torch.nn as nn +import torch.nn.functional as F +import torch +import numpy as np + +from ctlearn.core.pytorch.nets.models.EffientNet_pytorch.model import EfficientNet + +class MemoryEfficientSwish(nn.Module): + """ + Memory-efficient implementation of Swish activation function. + + Swish (also known as SiLU - Sigmoid Linear Unit) is a smooth, non-monotonic + activation function defined as: f(x) = x · σ(x) where σ is the sigmoid function. + + This implementation is memory-efficient because it doesn't store intermediate + values during the forward pass, reducing memory consumption during backpropagation. + + Mathematical Formula: + Swish(x) = x * sigmoid(x) = x * (1 / (1 + exp(-x))) + + Properties: + - Smooth and continuously differentiable + - Non-monotonic (can decrease for negative inputs) + - Bounded below (approaches 0 for large negative x) + - Unbounded above (approaches x for large positive x) + - Self-gated: combines input with its sigmoid + + Example: + >>> swish = MemoryEfficientSwish() + >>> x = torch.randn(32, 128) + >>> output = swish(x) + """ + + def forward(self, x): + """ + Apply Swish activation element-wise. + + Args: + x (torch.Tensor): Input tensor of any shape + + Returns: + torch.Tensor: Activated tensor with same shape as input + """ + return x * torch.sigmoid(x) + +class SEBlock(nn.Module): + """ + Squeeze-and-Excitation (SE) block for channel attention. + + This module implements channel-wise attention that adaptively recalibrates + channel-wise feature responses by explicitly modeling interdependencies + between channels. It "squeezes" global spatial information into a channel + descriptor and "excites" channels through a gating mechanism. + + Architecture: + Input → Global Avg Pool → FC (reduce) → PReLU → FC (expand) → Swish → Scale Input + + The SE block improves representational power by allowing the network to + emphasize informative features and suppress less useful ones. + + Attributes: + avg_pool (nn.AdaptiveAvgPool2d): Global average pooling layer + fc (nn.Sequential): Two-layer MLP for channel attention + - First layer: Channel reduction (compression) + - Second layer: Channel expansion (excitation) + """ + + def __init__(self, channel, reduction=16): + """ + Initialize the Squeeze-and-Excitation block. + + Args: + channel (int): Number of input channels + reduction (int, optional): Reduction ratio for the bottleneck. + Defaults to 16. Higher values reduce parameters but may lose information. + Common values: 4, 8, 16 + + Example: + >>> se_block = SEBlock(channel=512, reduction=16) + >>> x = torch.randn(8, 512, 7, 7) + >>> out = se_block(x) # Same shape as input + """ + super(SEBlock, self).__init__() + + # Global average pooling: (B, C, H, W) → (B, C, 1, 1) + self.avg_pool = nn.AdaptiveAvgPool2d(1) + + # Two-layer MLP for channel attention + self.fc = nn.Sequential( + # Squeeze: Reduce channel dimension + nn.Linear(channel, channel // reduction, bias=False), + nn.PReLU(), # Parametric ReLU activation + # Excitation: Restore channel dimension + nn.Linear(channel // reduction, channel, bias=False), + MemoryEfficientSwish() # Final activation for attention weights + ) + + def forward(self, x): + """ + Apply channel attention to input feature map. + + Process: + 1. Global average pool to get channel statistics (B, C, H, W) → (B, C) + 2. Pass through MLP to get channel attention weights + 3. Reshape weights to (B, C, 1, 1) + 4. Scale input features by attention weights (element-wise multiplication) + + Args: + x (torch.Tensor): Input feature map with shape (B, C, H, W) + B: batch size, C: channels, H: height, W: width + + Returns: + torch.Tensor: Attention-weighted feature map with same shape as input + Important channels are emphasized, less important ones suppressed + """ + b, c, _, _ = x.size() + + # Squeeze: Global average pooling + y = self.avg_pool(x).view(b, c) # (B, C, 1, 1) → (B, C) + + # Excitation: Learn channel attention weights + y = self.fc(y).view(b, c, 1, 1) # (B, C) → (B, C, 1, 1) + + # Scale input by attention weights + return x * y.expand_as(x) + +class DoubleBBEfficientNet(nn.Module): + """ + Dual-Backbone EfficientNet for multi-task telescope data analysis. + + This architecture uses two pretrained EfficientNet backbones to process + different input modalities (charge and timing images) independently, + then fuses their features for task-specific predictions. Supports multiple + EfficientNet variants (B0-B7) and different tasks (classification, regression). + + Architecture Overview: + Input_1 (charge) → EfficientNet_1 ↘ + → Fusion → Task-specific Head → Output + Input_2 (timing) → EfficientNet_2 ↗ + + Fusion Strategy: + - Addition fusion for computational efficiency + - 1x1 convolution for feature refinement + - Global average pooling for spatial reduction + + Task-Specific Heads: + - Classification: Two-layer MLP with Swish activation + - Energy: Dual-head (classification + regression) + - Direction: Two-layer MLP with dropout + + Attributes: + task (str): Task type ('type', 'energy', 'direction') + num_outputs (int): Number of output values + backbone1, backbone2 (EfficientNet): Feature extraction backbones + fusion_conv (nn.Conv2d): 1x1 conv for feature fusion + Task-specific layers (fc_*, prelu_*, dropout_*, etc.) + """ + + def __init__(self, model_variant: str = "efficientnet-b3", task: str = "Energy", + num_outputs=2, device_str="cuda", energy_bins=None): + """ + Initialize the Dual-Backbone EfficientNet model. + + Args: + model_variant (str, optional): EfficientNet variant to use. + Defaults to "efficientnet-b3" + Options: 'efficientnet-b0' to 'efficientnet-b7' + Larger models (b5, b7) have more parameters and accuracy + + task (str, optional): Task type. Defaults to "Energy" + Options: 'type' (classification), 'energy' (regression), 'direction' + + num_outputs (int, optional): Number of output values. Defaults to 2 + For 'type': 2 (gamma vs proton) + For 'energy': Variable (classification bins + regression) + For 'direction': 2 or 3 (angular coordinates) + + device_str (str, optional): Device string. Defaults to "cuda" + + energy_bins (list or None, optional): Energy bin edges for classification. + Defaults to None. Used only for energy task with dual-head approach. + + Raises: + ValueError: If model_variant is not tested/supported + + Example: + >>> # Classification with EfficientNet-B3 + >>> model = DoubleBBEfficientNet( + ... model_variant='efficientnet-b3', + ... task='type', + ... num_outputs=2 + ... ) + + >>> # Energy regression with larger model + >>> model = DoubleBBEfficientNet( + ... model_variant='efficientnet-b5', + ... task='energy', + ... num_outputs=10 + ... ) + """ + super(DoubleBBEfficientNet, self).__init__() + + # Store configuration + self.task = task.lower() + self.num_outputs = num_outputs + self.energy_bins = energy_bins + self.device = torch.device(device_str) + + # Define architecture parameters based on model variant + hidden_size = 512 + if 'b3' in model_variant: + feature_size = 1536 # EfficientNet-B3 feature dimension + elif 'b5' in model_variant: + feature_size = 2048 # EfficientNet-B5 feature dimension + else: + raise ValueError(f"Model variant {model_variant} not tested. Adapt the feature_size.") + + # Initialize backbones based on task + if self.task == "type": + # Classification task: use batch normalization + self.backbone1 = EfficientNet.from_pretrained( + model_variant, in_channels=1, num_classes=num_outputs, use_batch_norm=True + ) + self.backbone2 = EfficientNet.from_pretrained( + model_variant, in_channels=1, num_classes=num_outputs, use_batch_norm=True + ) + + elif self.task == "energy": + # Energy regression task + self.backbone1 = EfficientNet.from_pretrained( + model_variant, in_channels=1, num_classes=num_outputs + ) + self.backbone2 = EfficientNet.from_pretrained( + model_variant, in_channels=1, num_classes=num_outputs + ) + + elif self.task == "direction": + # Direction reconstruction task + self.backbone1 = EfficientNet.from_pretrained( + model_variant, in_channels=1, num_classes=num_outputs + ) + self.backbone2 = EfficientNet.from_pretrained( + model_variant, in_channels=1, num_classes=num_outputs + ) + + # Fusion module: 1x1 convolution to refine fused features + self.fusion_conv = nn.Conv2d( + in_channels=feature_size, + out_channels=feature_size, + kernel_size=1 + ) + + # Task-specific heads + if self.task == "type": + # Classification head: two-layer MLP + self.fc_classification_1 = nn.Linear(feature_size, hidden_size) + self.fc_classification_2 = nn.Linear(hidden_size, num_outputs) + + if self.task == "energy": + # Dual-head architecture for energy prediction + # Head 1: Energy range classification + self.fc_energy_1 = nn.Linear(feature_size, hidden_size) + self.fc_energy_2 = nn.Linear(hidden_size, int(num_outputs - 1)) + + # Head 2: Fine-grained regression + self.fc_energy_1_reg = nn.Linear(feature_size, hidden_size) + self.fc_energy_2_reg = nn.Linear(hidden_size, 1) + + if self.task == "direction": + # Direction head with dropout + self.fc_direction_1 = nn.Linear(feature_size, hidden_size) + self.fc_direction_2 = nn.Linear(hidden_size, num_outputs) + + # Activation and regularization layers + self.swish = MemoryEfficientSwish() + self.batch_norm = nn.BatchNorm1d(feature_size) + + # Task-specific dropout + self.dropout_energy_1 = nn.Dropout(0.1) + self.dropout_energy_2 = nn.Dropout(0.3) + self.dropout_direction = nn.Dropout(0.3) + + # Parametric activations + self.prelu_direction = nn.PReLU(num_parameters=hidden_size) + self.prelu_energy_1 = nn.PReLU(num_parameters=hidden_size) + self.prelu_energy_2 = nn.PReLU() + + def extract_feature_vector(self, x1, x2): + """ + Extract and fuse features from both backbones. + + This method processes both inputs through separate EfficientNet backbones, + fuses the features, and produces a compact feature vector for final prediction. + + Process: + 1. Extract features from both inputs using separate backbones + 2. Fuse features using element-wise addition + 3. Refine fused features with 1x1 convolution + 4. Apply global average pooling + 5. Flatten to feature vector + + Args: + x1 (torch.Tensor): First input (charge image) + Shape: (batch_size, 1, height, width) + x2 (torch.Tensor): Second input (timing image) + Shape: (batch_size, 1, height, width) + + Returns: + torch.Tensor: Fused feature vector + Shape: (batch_size, feature_size) + + Note: + Addition fusion is used for computational efficiency. + Alternative fusion strategies (concatenation, weighted sum) are possible. + """ + # Extract features from both backbones + x1 = self.backbone1.extract_features(x1) + x2 = self.backbone2.extract_features(x2) + + # Fusion: Element-wise addition + # Alternative: torch.cat((x1, x2), dim=1) for concatenation + fused_features = torch.add(x1, x2) + + # Refine fused features with 1x1 convolution + fused_features = self.fusion_conv(fused_features) + + # Global average pooling: (B, C, H, W) → (B, C, 1, 1) + fused_features = F.adaptive_avg_pool2d(fused_features, 1) + + # Flatten: (B, C, 1, 1) → (B, C) + fused_features = fused_features.view(fused_features.size(0), -1) + + return fused_features + + def forward(self, x1): + if x1.shape[1] >= 2: + x1, x2 = torch.split(x1, [1, x1.shape[1]-1], dim=1) + else: + x2 = x1 + """ + Forward pass through the dual-backbone network. + + Process: + 1. Extract and fuse features from both inputs + 2. Apply normalization and dropout (for energy/direction) + 3. Pass through task-specific prediction head + 4. Return predictions in standardized format + + Args: + x1 (torch.Tensor): First input (charge image) + Shape: (batch_size, 1, height, width) + x2 (torch.Tensor): Second input (timing image) + Shape: (batch_size, 1, height, width) + + Returns: + tuple: (classification, energy, direction) where: + - classification: For 'type' task + [logits, features] where logits: (batch_size, 2) + - energy: For 'energy' task + [class_logits, regression_value] + class_logits: (batch_size, num_bins-1) + regression_value: (batch_size, 1) + - direction: For 'direction' task + [predictions, features] where predictions: (batch_size, num_outputs) + + Only one of the three is non-None based on self.task + + Task-Specific Processing: + Classification (type): + - Two-layer MLP with Swish activation + - Returns logits and feature vector + + Energy (energy): + - Dual-head architecture + - Classification head: Predicts energy range/bin + - Regression head: Predicts fine-grained value + - Combines coarse and fine predictions + + Direction (direction): + - Two-layer MLP with PReLU and dropout + - Predicts angular offsets or coordinates + """ + # Initialize outputs + energy = [None, None] + classification = [None, None] + direction = [None, None] + + # Extract fused features from both backbones + fused_features = self.extract_feature_vector(x1, x2) + + # Task-specific prediction heads + if self.task == "type": + # Classification head + classification = self.fc_classification_1(fused_features) + classification = self.fc_classification_2(classification) + classification = self.swish(classification) + + if self.task == "energy": + # Apply normalization and dropout + fused_features = self.batch_norm(fused_features) + fused_features = self.dropout_energy_1(fused_features) + + # Classification head: Predict energy bin + energy_class = self.fc_energy_1(fused_features) + energy_class = self.fc_energy_2(energy_class) + energy_class = self.swish(energy_class) + energy_pred_class = self.swish(energy_class) + + # Regression head: Predict fine-grained energy + energy_reg = self.fc_energy_1_reg(fused_features) + energy_reg = self.dropout_energy_2(energy_reg) + energy_reg = self.prelu_energy_1(energy_reg) + energy_regresion = self.fc_energy_2_reg(energy_reg) + + # Combine both predictions + energy = [energy_pred_class, energy_regresion] + + if self.task == "direction": + # Direction head with dropout + fused_features = self.dropout_direction(fused_features) + direction = self.fc_direction_1(fused_features) + direction = self.fc_direction_2(direction) + + # Return in standardized format + return [classification, fused_features], energy, direction + + def eval(self): + """ + Set the model to evaluation mode. + + Overrides the default eval() to ensure both backbones are also + set to evaluation mode. This is important for proper handling of + batch normalization and dropout layers. + """ + super().eval() + self.backbone1.eval() + self.backbone2.eval() + + def train(self, mode=True): + """ + Set the model to training or evaluation mode. + + Overrides the default train() to ensure both backbones follow + the same mode. This ensures consistent behavior of batch normalization + and dropout across all components. + + Args: + mode (bool, optional): Whether to set training mode (True) or + evaluation mode (False). Defaults to True. + """ + super().train(mode) + self.backbone1.train(mode) + self.backbone2.train(mode) diff --git a/ctlearn/core/pytorch/nets/models/DoubleBBEfficientNet/__init__.py b/ctlearn/core/pytorch/nets/models/DoubleBBEfficientNet/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/__init__.py b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/__init__.py new file mode 100644 index 00000000..2b529dfe --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/__init__.py @@ -0,0 +1,9 @@ +__version__ = "0.7.1" +from .model import EfficientNet, VALID_MODELS +from .utils import ( + GlobalParams, + BlockArgs, + BlockDecoder, + efficientnet, + get_model_params, +) diff --git a/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/model.py b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/model.py new file mode 100644 index 00000000..97816df8 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/model.py @@ -0,0 +1,454 @@ +"""model.py - Model and module class for EfficientNet. + They are built to mirror those in the official TensorFlow implementation. +""" + +# Author: lukemelas (github username) +# Github repo: https://github.com/lukemelas/EfficientNet-PyTorch +# With adjustments and added comments by workingcoder (github username). + +import torch +from torch import nn +from torch.nn import functional as F +from .utils import ( + round_filters, + round_repeats, + drop_connect, + get_same_padding_conv2d, + get_model_params, + efficientnet_params, + load_pretrained_weights, + Swish, + MemoryEfficientSwish, + calculate_output_image_size +) + + +VALID_MODELS = ( + 'efficientnet-b0', 'efficientnet-b1', 'efficientnet-b2', 'efficientnet-b3', + 'efficientnet-b4', 'efficientnet-b5', 'efficientnet-b6', 'efficientnet-b7', + 'efficientnet-b8', + + # Support the construction of 'efficientnet-l2' without pretrained weights + 'efficientnet-l2' +) + + +class MBConvBlock(nn.Module): + """Mobile Inverted Residual Bottleneck Block. + + Args: + block_args (namedtuple): BlockArgs, defined in utils.py. + global_params (namedtuple): GlobalParam, defined in utils.py. + image_size (tuple or list): [image_height, image_width]. + + References: + [1] https://arxiv.org/abs/1704.04861 (MobileNet v1) + [2] https://arxiv.org/abs/1801.04381 (MobileNet v2) + [3] https://arxiv.org/abs/1905.02244 (MobileNet v3) + """ + + def __init__(self, block_args, global_params, image_size=None,use_swish=True, use_batch_norm=True): + super().__init__() + self._block_args = block_args + self._bn_mom = 1 - global_params.batch_norm_momentum # pytorch's difference from tensorflow + self._bn_eps = global_params.batch_norm_epsilon + self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1) + self.id_skip = block_args.id_skip # whether to use skip connection and drop connect + self.use_swish = use_swish + self.use_batch_norm = use_batch_norm + # Expansion phase (Inverted Bottleneck) + inp = self._block_args.input_filters # number of input channels + oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels + if self._block_args.expand_ratio != 1: + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False) + if self.use_batch_norm: + self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) + else: + self._bn0 = nn.Identity() + + # image_size = calculate_output_image_size(image_size, 1) <-- this wouldn't modify image_size + + # Depthwise convolution phase + k = self._block_args.kernel_size + s = self._block_args.stride + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._depthwise_conv = Conv2d( + in_channels=oup, out_channels=oup, groups=oup, # groups makes it depthwise + kernel_size=k, stride=s, bias=False) + if self.use_batch_norm: + self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) + else: + self._bn1 = nn.Identity() + + image_size = calculate_output_image_size(image_size, s) + + # Squeeze and Excitation layer, if desired + if self.has_se: + Conv2d = get_same_padding_conv2d(image_size=(1, 1)) + num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio)) + self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1) + self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1) + + # Pointwise convolution phase + final_oup = self._block_args.output_filters + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False) + if self.use_batch_norm: + self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps) + else: + self._bn2 = nn.Identity() + if self.use_swish: + self._swish = MemoryEfficientSwish() + else: + self._swish = nn.PReLU() + + def forward(self, inputs, drop_connect_rate=None): + """MBConvBlock's forward function. + + Args: + inputs (tensor): Input tensor. + drop_connect_rate (bool): Drop connect rate (float, between 0 and 1). + + Returns: + Output of this block after processing. + """ + + # Expansion and Depthwise Convolution + x = inputs + if self._block_args.expand_ratio != 1: + x = self._expand_conv(inputs) + x = self._bn0(x) + x = self._swish(x) + + x = self._depthwise_conv(x) + x = self._bn1(x) + x = self._swish(x) + + # Squeeze and Excitation + if self.has_se: + x_squeezed = F.adaptive_avg_pool2d(x, 1) + x_squeezed = self._se_reduce(x_squeezed) + x_squeezed = self._swish(x_squeezed) + x_squeezed = self._se_expand(x_squeezed) + x = torch.sigmoid(x_squeezed) * x + + # Pointwise Convolution + x = self._project_conv(x) + x = self._bn2(x) + + # Skip connection and drop connect + input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters + if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters: + # The combination of skip connection and drop connect brings about stochastic depth. + if drop_connect_rate: + x = drop_connect(x, p=drop_connect_rate, training=self.training) + x = x + inputs # skip connection + return x + + def set_swish(self, memory_efficient=True): + """Sets swish function as memory efficient (for training) or standard (for export). + + Args: + memory_efficient (bool): Whether to use memory-efficient version of swish. + """ + self._swish = MemoryEfficientSwish() if memory_efficient else Swish() + + +class EfficientNet(nn.Module): + """EfficientNet model. + Most easily loaded with the .from_name or .from_pretrained methods. + + Args: + blocks_args (list[namedtuple]): A list of BlockArgs to construct blocks. + global_params (namedtuple): A set of GlobalParams shared between blocks. + + References: + [1] https://arxiv.org/abs/1905.11946 (EfficientNet) + + Example: + >>> import torch + >>> from efficientnet.model import EfficientNet + >>> inputs = torch.rand(1, 3, 224, 224) + >>> model = EfficientNet.from_pretrained('efficientnet-b0') + >>> model.eval() + >>> outputs = model(inputs) + """ + + def __init__(self, blocks_args=None, global_params=None,use_swish=True,use_bn=True): + + # DELETE MEEEEEEE + # use_bn=False + + super().__init__() + assert isinstance(blocks_args, list), 'blocks_args should be a list' + assert len(blocks_args) > 0, 'block args must be greater than 0' + self._global_params = global_params + self._blocks_args = blocks_args + self.use_swish= use_swish + self.use_batch_norm = use_bn + # Batch norm parameters + bn_mom = 1 - self._global_params.batch_norm_momentum + bn_eps = self._global_params.batch_norm_epsilon + + # Get stem static or dynamic convolution depending on image size + image_size = global_params.image_size + Conv2d = get_same_padding_conv2d(image_size=image_size) + + # Stem + in_channels = 3 # rgb + out_channels = round_filters(32, self._global_params) # number of output channels + self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False) + if self.use_batch_norm: + self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) + else: + self._bn0 = nn.Identity() + image_size = calculate_output_image_size(image_size, 2) + + # Build blocks + self._blocks = nn.ModuleList([]) + for block_args in self._blocks_args: + + # Update block input and output filters based on depth multiplier. + block_args = block_args._replace( + input_filters=round_filters(block_args.input_filters, self._global_params), + output_filters=round_filters(block_args.output_filters, self._global_params), + num_repeat=round_repeats(block_args.num_repeat, self._global_params) + ) + + # The first block needs to take care of stride and filter size increase. + self._blocks.append(MBConvBlock(block_args, self._global_params, image_size=image_size,use_swish=self.use_swish, use_batch_norm=self.use_batch_norm)) + image_size = calculate_output_image_size(image_size, block_args.stride) + if block_args.num_repeat > 1: # modify block_args to keep same output size + block_args = block_args._replace(input_filters=block_args.output_filters, stride=1) + for _ in range(block_args.num_repeat - 1): + self._blocks.append(MBConvBlock(block_args, self._global_params, image_size=image_size,use_swish=self.use_swish,use_batch_norm=self.use_batch_norm)) + # image_size = calculate_output_image_size(image_size, block_args.stride) # stride = 1 + + # Head + in_channels = block_args.output_filters # output of final block + out_channels = round_filters(1280, self._global_params) + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + if self.use_batch_norm: + self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) + else: + self._bn1 = nn.Identity() + + # Final linear layer + self._avg_pooling = nn.AdaptiveAvgPool2d(1) + if self._global_params.include_top: + self._dropout = nn.Dropout(self._global_params.dropout_rate) + self._fc = nn.Linear(out_channels, self._global_params.num_classes) + + # set activation to memory efficient swish by default + if self.use_swish: + self._swish = MemoryEfficientSwish() + else: + self._swish = nn.PReLU() + + def set_swish(self, memory_efficient=True): + """Sets swish function as memory efficient (for training) or standard (for export). + + Args: + memory_efficient (bool): Whether to use memory-efficient version of swish. + """ + if self.use_swish: + self._swish = MemoryEfficientSwish() if memory_efficient else Swish() + else: + self._swish = nn.PReLU() + + for block in self._blocks: + block.set_swish(memory_efficient) + + def extract_endpoints(self, inputs): + """Use convolution layer to extract features + from reduction levels i in [1, 2, 3, 4, 5]. + + Args: + inputs (tensor): Input tensor. + + Returns: + Dictionary of last intermediate features + with reduction levels i in [1, 2, 3, 4, 5]. + Example: + >>> import torch + >>> from efficientnet.model import EfficientNet + >>> inputs = torch.rand(1, 3, 224, 224) + >>> model = EfficientNet.from_pretrained('efficientnet-b0') + >>> endpoints = model.extract_endpoints(inputs) + >>> print(endpoints['reduction_1'].shape) # torch.Size([1, 16, 112, 112]) + >>> print(endpoints['reduction_2'].shape) # torch.Size([1, 24, 56, 56]) + >>> print(endpoints['reduction_3'].shape) # torch.Size([1, 40, 28, 28]) + >>> print(endpoints['reduction_4'].shape) # torch.Size([1, 112, 14, 14]) + >>> print(endpoints['reduction_5'].shape) # torch.Size([1, 320, 7, 7]) + >>> print(endpoints['reduction_6'].shape) # torch.Size([1, 1280, 7, 7]) + """ + endpoints = dict() + + # Stem + x = self._swish(self._bn0(self._conv_stem(inputs))) + prev_x = x + + # Blocks + for idx, block in enumerate(self._blocks): + drop_connect_rate = self._global_params.drop_connect_rate + if drop_connect_rate: + drop_connect_rate *= float(idx) / len(self._blocks) # scale drop connect_rate + x = block(x, drop_connect_rate=drop_connect_rate) + if prev_x.size(2) > x.size(2): + endpoints['reduction_{}'.format(len(endpoints) + 1)] = prev_x + elif idx == len(self._blocks) - 1: + endpoints['reduction_{}'.format(len(endpoints) + 1)] = x + prev_x = x + + # Head + x = self._swish(self._bn1(self._conv_head(x))) + endpoints['reduction_{}'.format(len(endpoints) + 1)] = x + + return endpoints + + def extract_features(self, inputs): + """use convolution layer to extract feature . + + Args: + inputs (tensor): Input tensor. + + Returns: + Output of the final convolution + layer in the efficientnet model. + """ + # Stem + x = self._swish(self._bn0(self._conv_stem(inputs))) + + # Blocks + for idx, block in enumerate(self._blocks): + drop_connect_rate = self._global_params.drop_connect_rate + drop_connect_rate=0.0 + if drop_connect_rate: + drop_connect_rate *= float(idx) / len(self._blocks) # scale drop connect_rate + x = block(x, drop_connect_rate=drop_connect_rate) + + # Head + x = self._swish(self._bn1(self._conv_head(x))) + + return x + + def forward(self, inputs): + """EfficientNet's forward function. + Calls extract_features to extract features, applies final linear layer, and returns logits. + + Args: + inputs (tensor): Input tensor. + + Returns: + Output of this model after processing. + """ + # Convolution layers + x = self.extract_features(inputs) + # Pooling and final linear layer + x = self._avg_pooling(x) + if self._global_params.include_top: + x = x.flatten(start_dim=1) + x = self._dropout(x) + x = self._fc(x) + return x + + @classmethod + def from_name(cls, model_name, in_channels=3,use_swish=True,use_batch_norm=True, **override_params): + """Create an efficientnet model according to name. + + Args: + model_name (str): Name for efficientnet. + in_channels (int): Input data's channel number. + override_params (other key word params): + Params to override model's global_params. + Optional key: + 'width_coefficient', 'depth_coefficient', + 'image_size', 'dropout_rate', + 'num_classes', 'batch_norm_momentum', + 'batch_norm_epsilon', 'drop_connect_rate', + 'depth_divisor', 'min_depth' + + Returns: + An efficientnet model. + """ + cls._check_model_name_is_valid(model_name) + blocks_args, global_params = get_model_params(model_name, override_params) + model = cls(blocks_args, global_params,use_swish,use_batch_norm) + model._change_in_channels(in_channels) + + return model + + @classmethod + def from_pretrained(cls, model_name, weights_path=None, advprop=False, + in_channels=3, num_classes=1000, use_batch_norm=True, **override_params): + """Create an efficientnet model according to name. + + Args: + model_name (str): Name for efficientnet. + weights_path (None or str): + str: path to pretrained weights file on the local disk. + None: use pretrained weights downloaded from the Internet. + advprop (bool): + Whether to load pretrained weights + trained with advprop (valid when weights_path is None). + in_channels (int): Input data's channel number. + num_classes (int): + Number of categories for classification. + It controls the output size for final linear layer. + override_params (other key word params): + Params to override model's global_params. + Optional key: + 'width_coefficient', 'depth_coefficient', + 'image_size', 'dropout_rate', + 'batch_norm_momentum', + 'batch_norm_epsilon', 'drop_connect_rate', + 'depth_divisor', 'min_depth' + + Returns: + A pretrained efficientnet model. + """ + model = cls.from_name(model_name, num_classes=num_classes,use_batch_norm=use_batch_norm, **override_params) + # load_pretrained_weights(model, model_name, weights_path=weights_path,load_fc=(num_classes == 1000), advprop=advprop) + model._change_in_channels(in_channels) + + return model + + @classmethod + def get_image_size(cls, model_name): + """Get the input image size for a given efficientnet model. + + Args: + model_name (str): Name for efficientnet. + + Returns: + Input image size (resolution). + """ + cls._check_model_name_is_valid(model_name) + _, _, res, _ = efficientnet_params(model_name) + return res + + @classmethod + def _check_model_name_is_valid(cls, model_name): + """Validates model name. + + Args: + model_name (str): Name for efficientnet. + + Returns: + bool: Is a valid name or not. + """ + if model_name not in VALID_MODELS: + raise ValueError('model_name should be one of: ' + ', '.join(VALID_MODELS)) + + def _change_in_channels(self, in_channels): + """Adjust model's first convolution layer to in_channels, if in_channels not equals 3. + + Args: + in_channels (int): Input data's channel number. + """ + if in_channels != 3: + Conv2d = get_same_padding_conv2d(image_size=self._global_params.image_size) + out_channels = round_filters(32, self._global_params) + self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False) diff --git a/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/model_original.py b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/model_original.py new file mode 100644 index 00000000..97816df8 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/model_original.py @@ -0,0 +1,454 @@ +"""model.py - Model and module class for EfficientNet. + They are built to mirror those in the official TensorFlow implementation. +""" + +# Author: lukemelas (github username) +# Github repo: https://github.com/lukemelas/EfficientNet-PyTorch +# With adjustments and added comments by workingcoder (github username). + +import torch +from torch import nn +from torch.nn import functional as F +from .utils import ( + round_filters, + round_repeats, + drop_connect, + get_same_padding_conv2d, + get_model_params, + efficientnet_params, + load_pretrained_weights, + Swish, + MemoryEfficientSwish, + calculate_output_image_size +) + + +VALID_MODELS = ( + 'efficientnet-b0', 'efficientnet-b1', 'efficientnet-b2', 'efficientnet-b3', + 'efficientnet-b4', 'efficientnet-b5', 'efficientnet-b6', 'efficientnet-b7', + 'efficientnet-b8', + + # Support the construction of 'efficientnet-l2' without pretrained weights + 'efficientnet-l2' +) + + +class MBConvBlock(nn.Module): + """Mobile Inverted Residual Bottleneck Block. + + Args: + block_args (namedtuple): BlockArgs, defined in utils.py. + global_params (namedtuple): GlobalParam, defined in utils.py. + image_size (tuple or list): [image_height, image_width]. + + References: + [1] https://arxiv.org/abs/1704.04861 (MobileNet v1) + [2] https://arxiv.org/abs/1801.04381 (MobileNet v2) + [3] https://arxiv.org/abs/1905.02244 (MobileNet v3) + """ + + def __init__(self, block_args, global_params, image_size=None,use_swish=True, use_batch_norm=True): + super().__init__() + self._block_args = block_args + self._bn_mom = 1 - global_params.batch_norm_momentum # pytorch's difference from tensorflow + self._bn_eps = global_params.batch_norm_epsilon + self.has_se = (self._block_args.se_ratio is not None) and (0 < self._block_args.se_ratio <= 1) + self.id_skip = block_args.id_skip # whether to use skip connection and drop connect + self.use_swish = use_swish + self.use_batch_norm = use_batch_norm + # Expansion phase (Inverted Bottleneck) + inp = self._block_args.input_filters # number of input channels + oup = self._block_args.input_filters * self._block_args.expand_ratio # number of output channels + if self._block_args.expand_ratio != 1: + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._expand_conv = Conv2d(in_channels=inp, out_channels=oup, kernel_size=1, bias=False) + if self.use_batch_norm: + self._bn0 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) + else: + self._bn0 = nn.Identity() + + # image_size = calculate_output_image_size(image_size, 1) <-- this wouldn't modify image_size + + # Depthwise convolution phase + k = self._block_args.kernel_size + s = self._block_args.stride + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._depthwise_conv = Conv2d( + in_channels=oup, out_channels=oup, groups=oup, # groups makes it depthwise + kernel_size=k, stride=s, bias=False) + if self.use_batch_norm: + self._bn1 = nn.BatchNorm2d(num_features=oup, momentum=self._bn_mom, eps=self._bn_eps) + else: + self._bn1 = nn.Identity() + + image_size = calculate_output_image_size(image_size, s) + + # Squeeze and Excitation layer, if desired + if self.has_se: + Conv2d = get_same_padding_conv2d(image_size=(1, 1)) + num_squeezed_channels = max(1, int(self._block_args.input_filters * self._block_args.se_ratio)) + self._se_reduce = Conv2d(in_channels=oup, out_channels=num_squeezed_channels, kernel_size=1) + self._se_expand = Conv2d(in_channels=num_squeezed_channels, out_channels=oup, kernel_size=1) + + # Pointwise convolution phase + final_oup = self._block_args.output_filters + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._project_conv = Conv2d(in_channels=oup, out_channels=final_oup, kernel_size=1, bias=False) + if self.use_batch_norm: + self._bn2 = nn.BatchNorm2d(num_features=final_oup, momentum=self._bn_mom, eps=self._bn_eps) + else: + self._bn2 = nn.Identity() + if self.use_swish: + self._swish = MemoryEfficientSwish() + else: + self._swish = nn.PReLU() + + def forward(self, inputs, drop_connect_rate=None): + """MBConvBlock's forward function. + + Args: + inputs (tensor): Input tensor. + drop_connect_rate (bool): Drop connect rate (float, between 0 and 1). + + Returns: + Output of this block after processing. + """ + + # Expansion and Depthwise Convolution + x = inputs + if self._block_args.expand_ratio != 1: + x = self._expand_conv(inputs) + x = self._bn0(x) + x = self._swish(x) + + x = self._depthwise_conv(x) + x = self._bn1(x) + x = self._swish(x) + + # Squeeze and Excitation + if self.has_se: + x_squeezed = F.adaptive_avg_pool2d(x, 1) + x_squeezed = self._se_reduce(x_squeezed) + x_squeezed = self._swish(x_squeezed) + x_squeezed = self._se_expand(x_squeezed) + x = torch.sigmoid(x_squeezed) * x + + # Pointwise Convolution + x = self._project_conv(x) + x = self._bn2(x) + + # Skip connection and drop connect + input_filters, output_filters = self._block_args.input_filters, self._block_args.output_filters + if self.id_skip and self._block_args.stride == 1 and input_filters == output_filters: + # The combination of skip connection and drop connect brings about stochastic depth. + if drop_connect_rate: + x = drop_connect(x, p=drop_connect_rate, training=self.training) + x = x + inputs # skip connection + return x + + def set_swish(self, memory_efficient=True): + """Sets swish function as memory efficient (for training) or standard (for export). + + Args: + memory_efficient (bool): Whether to use memory-efficient version of swish. + """ + self._swish = MemoryEfficientSwish() if memory_efficient else Swish() + + +class EfficientNet(nn.Module): + """EfficientNet model. + Most easily loaded with the .from_name or .from_pretrained methods. + + Args: + blocks_args (list[namedtuple]): A list of BlockArgs to construct blocks. + global_params (namedtuple): A set of GlobalParams shared between blocks. + + References: + [1] https://arxiv.org/abs/1905.11946 (EfficientNet) + + Example: + >>> import torch + >>> from efficientnet.model import EfficientNet + >>> inputs = torch.rand(1, 3, 224, 224) + >>> model = EfficientNet.from_pretrained('efficientnet-b0') + >>> model.eval() + >>> outputs = model(inputs) + """ + + def __init__(self, blocks_args=None, global_params=None,use_swish=True,use_bn=True): + + # DELETE MEEEEEEE + # use_bn=False + + super().__init__() + assert isinstance(blocks_args, list), 'blocks_args should be a list' + assert len(blocks_args) > 0, 'block args must be greater than 0' + self._global_params = global_params + self._blocks_args = blocks_args + self.use_swish= use_swish + self.use_batch_norm = use_bn + # Batch norm parameters + bn_mom = 1 - self._global_params.batch_norm_momentum + bn_eps = self._global_params.batch_norm_epsilon + + # Get stem static or dynamic convolution depending on image size + image_size = global_params.image_size + Conv2d = get_same_padding_conv2d(image_size=image_size) + + # Stem + in_channels = 3 # rgb + out_channels = round_filters(32, self._global_params) # number of output channels + self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False) + if self.use_batch_norm: + self._bn0 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) + else: + self._bn0 = nn.Identity() + image_size = calculate_output_image_size(image_size, 2) + + # Build blocks + self._blocks = nn.ModuleList([]) + for block_args in self._blocks_args: + + # Update block input and output filters based on depth multiplier. + block_args = block_args._replace( + input_filters=round_filters(block_args.input_filters, self._global_params), + output_filters=round_filters(block_args.output_filters, self._global_params), + num_repeat=round_repeats(block_args.num_repeat, self._global_params) + ) + + # The first block needs to take care of stride and filter size increase. + self._blocks.append(MBConvBlock(block_args, self._global_params, image_size=image_size,use_swish=self.use_swish, use_batch_norm=self.use_batch_norm)) + image_size = calculate_output_image_size(image_size, block_args.stride) + if block_args.num_repeat > 1: # modify block_args to keep same output size + block_args = block_args._replace(input_filters=block_args.output_filters, stride=1) + for _ in range(block_args.num_repeat - 1): + self._blocks.append(MBConvBlock(block_args, self._global_params, image_size=image_size,use_swish=self.use_swish,use_batch_norm=self.use_batch_norm)) + # image_size = calculate_output_image_size(image_size, block_args.stride) # stride = 1 + + # Head + in_channels = block_args.output_filters # output of final block + out_channels = round_filters(1280, self._global_params) + Conv2d = get_same_padding_conv2d(image_size=image_size) + self._conv_head = Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + if self.use_batch_norm: + self._bn1 = nn.BatchNorm2d(num_features=out_channels, momentum=bn_mom, eps=bn_eps) + else: + self._bn1 = nn.Identity() + + # Final linear layer + self._avg_pooling = nn.AdaptiveAvgPool2d(1) + if self._global_params.include_top: + self._dropout = nn.Dropout(self._global_params.dropout_rate) + self._fc = nn.Linear(out_channels, self._global_params.num_classes) + + # set activation to memory efficient swish by default + if self.use_swish: + self._swish = MemoryEfficientSwish() + else: + self._swish = nn.PReLU() + + def set_swish(self, memory_efficient=True): + """Sets swish function as memory efficient (for training) or standard (for export). + + Args: + memory_efficient (bool): Whether to use memory-efficient version of swish. + """ + if self.use_swish: + self._swish = MemoryEfficientSwish() if memory_efficient else Swish() + else: + self._swish = nn.PReLU() + + for block in self._blocks: + block.set_swish(memory_efficient) + + def extract_endpoints(self, inputs): + """Use convolution layer to extract features + from reduction levels i in [1, 2, 3, 4, 5]. + + Args: + inputs (tensor): Input tensor. + + Returns: + Dictionary of last intermediate features + with reduction levels i in [1, 2, 3, 4, 5]. + Example: + >>> import torch + >>> from efficientnet.model import EfficientNet + >>> inputs = torch.rand(1, 3, 224, 224) + >>> model = EfficientNet.from_pretrained('efficientnet-b0') + >>> endpoints = model.extract_endpoints(inputs) + >>> print(endpoints['reduction_1'].shape) # torch.Size([1, 16, 112, 112]) + >>> print(endpoints['reduction_2'].shape) # torch.Size([1, 24, 56, 56]) + >>> print(endpoints['reduction_3'].shape) # torch.Size([1, 40, 28, 28]) + >>> print(endpoints['reduction_4'].shape) # torch.Size([1, 112, 14, 14]) + >>> print(endpoints['reduction_5'].shape) # torch.Size([1, 320, 7, 7]) + >>> print(endpoints['reduction_6'].shape) # torch.Size([1, 1280, 7, 7]) + """ + endpoints = dict() + + # Stem + x = self._swish(self._bn0(self._conv_stem(inputs))) + prev_x = x + + # Blocks + for idx, block in enumerate(self._blocks): + drop_connect_rate = self._global_params.drop_connect_rate + if drop_connect_rate: + drop_connect_rate *= float(idx) / len(self._blocks) # scale drop connect_rate + x = block(x, drop_connect_rate=drop_connect_rate) + if prev_x.size(2) > x.size(2): + endpoints['reduction_{}'.format(len(endpoints) + 1)] = prev_x + elif idx == len(self._blocks) - 1: + endpoints['reduction_{}'.format(len(endpoints) + 1)] = x + prev_x = x + + # Head + x = self._swish(self._bn1(self._conv_head(x))) + endpoints['reduction_{}'.format(len(endpoints) + 1)] = x + + return endpoints + + def extract_features(self, inputs): + """use convolution layer to extract feature . + + Args: + inputs (tensor): Input tensor. + + Returns: + Output of the final convolution + layer in the efficientnet model. + """ + # Stem + x = self._swish(self._bn0(self._conv_stem(inputs))) + + # Blocks + for idx, block in enumerate(self._blocks): + drop_connect_rate = self._global_params.drop_connect_rate + drop_connect_rate=0.0 + if drop_connect_rate: + drop_connect_rate *= float(idx) / len(self._blocks) # scale drop connect_rate + x = block(x, drop_connect_rate=drop_connect_rate) + + # Head + x = self._swish(self._bn1(self._conv_head(x))) + + return x + + def forward(self, inputs): + """EfficientNet's forward function. + Calls extract_features to extract features, applies final linear layer, and returns logits. + + Args: + inputs (tensor): Input tensor. + + Returns: + Output of this model after processing. + """ + # Convolution layers + x = self.extract_features(inputs) + # Pooling and final linear layer + x = self._avg_pooling(x) + if self._global_params.include_top: + x = x.flatten(start_dim=1) + x = self._dropout(x) + x = self._fc(x) + return x + + @classmethod + def from_name(cls, model_name, in_channels=3,use_swish=True,use_batch_norm=True, **override_params): + """Create an efficientnet model according to name. + + Args: + model_name (str): Name for efficientnet. + in_channels (int): Input data's channel number. + override_params (other key word params): + Params to override model's global_params. + Optional key: + 'width_coefficient', 'depth_coefficient', + 'image_size', 'dropout_rate', + 'num_classes', 'batch_norm_momentum', + 'batch_norm_epsilon', 'drop_connect_rate', + 'depth_divisor', 'min_depth' + + Returns: + An efficientnet model. + """ + cls._check_model_name_is_valid(model_name) + blocks_args, global_params = get_model_params(model_name, override_params) + model = cls(blocks_args, global_params,use_swish,use_batch_norm) + model._change_in_channels(in_channels) + + return model + + @classmethod + def from_pretrained(cls, model_name, weights_path=None, advprop=False, + in_channels=3, num_classes=1000, use_batch_norm=True, **override_params): + """Create an efficientnet model according to name. + + Args: + model_name (str): Name for efficientnet. + weights_path (None or str): + str: path to pretrained weights file on the local disk. + None: use pretrained weights downloaded from the Internet. + advprop (bool): + Whether to load pretrained weights + trained with advprop (valid when weights_path is None). + in_channels (int): Input data's channel number. + num_classes (int): + Number of categories for classification. + It controls the output size for final linear layer. + override_params (other key word params): + Params to override model's global_params. + Optional key: + 'width_coefficient', 'depth_coefficient', + 'image_size', 'dropout_rate', + 'batch_norm_momentum', + 'batch_norm_epsilon', 'drop_connect_rate', + 'depth_divisor', 'min_depth' + + Returns: + A pretrained efficientnet model. + """ + model = cls.from_name(model_name, num_classes=num_classes,use_batch_norm=use_batch_norm, **override_params) + # load_pretrained_weights(model, model_name, weights_path=weights_path,load_fc=(num_classes == 1000), advprop=advprop) + model._change_in_channels(in_channels) + + return model + + @classmethod + def get_image_size(cls, model_name): + """Get the input image size for a given efficientnet model. + + Args: + model_name (str): Name for efficientnet. + + Returns: + Input image size (resolution). + """ + cls._check_model_name_is_valid(model_name) + _, _, res, _ = efficientnet_params(model_name) + return res + + @classmethod + def _check_model_name_is_valid(cls, model_name): + """Validates model name. + + Args: + model_name (str): Name for efficientnet. + + Returns: + bool: Is a valid name or not. + """ + if model_name not in VALID_MODELS: + raise ValueError('model_name should be one of: ' + ', '.join(VALID_MODELS)) + + def _change_in_channels(self, in_channels): + """Adjust model's first convolution layer to in_channels, if in_channels not equals 3. + + Args: + in_channels (int): Input data's channel number. + """ + if in_channels != 3: + Conv2d = get_same_padding_conv2d(image_size=self._global_params.image_size) + out_channels = round_filters(32, self._global_params) + self._conv_stem = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, bias=False) diff --git a/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/utils.py b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/utils.py new file mode 100644 index 00000000..826a6279 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/EffientNet_pytorch/utils.py @@ -0,0 +1,616 @@ +"""utils.py - Helper functions for building the model and for loading model parameters. + These helper functions are built to mirror those in the official TensorFlow implementation. +""" + +# Author: lukemelas (github username) +# Github repo: https://github.com/lukemelas/EfficientNet-PyTorch +# With adjustments and added comments by workingcoder (github username). + +import re +import math +import collections +from functools import partial +import torch +from torch import nn +from torch.nn import functional as F +from torch.utils import model_zoo + + +################################################################################ +# Help functions for model architecture +################################################################################ + +# GlobalParams and BlockArgs: Two namedtuples +# Swish and MemoryEfficientSwish: Two implementations of the method +# round_filters and round_repeats: +# Functions to calculate params for scaling model width and depth ! ! ! +# get_width_and_height_from_size and calculate_output_image_size +# drop_connect: A structural design +# get_same_padding_conv2d: +# Conv2dDynamicSamePadding +# Conv2dStaticSamePadding +# get_same_padding_maxPool2d: +# MaxPool2dDynamicSamePadding +# MaxPool2dStaticSamePadding +# It's an additional function, not used in EfficientNet, +# but can be used in other model (such as EfficientDet). + +# Parameters for the entire model (stem, all blocks, and head) +GlobalParams = collections.namedtuple('GlobalParams', [ + 'width_coefficient', 'depth_coefficient', 'image_size', 'dropout_rate', + 'num_classes', 'batch_norm_momentum', 'batch_norm_epsilon', + 'drop_connect_rate', 'depth_divisor', 'min_depth', 'include_top']) + +# Parameters for an individual model block +BlockArgs = collections.namedtuple('BlockArgs', [ + 'num_repeat', 'kernel_size', 'stride', 'expand_ratio', + 'input_filters', 'output_filters', 'se_ratio', 'id_skip']) + +# Set GlobalParams and BlockArgs's defaults +GlobalParams.__new__.__defaults__ = (None,) * len(GlobalParams._fields) +BlockArgs.__new__.__defaults__ = (None,) * len(BlockArgs._fields) + +# Swish activation function +if hasattr(nn, 'SiLU'): + Swish = nn.SiLU +else: + # For compatibility with old PyTorch versions + class Swish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + + +# A memory-efficient implementation of Swish function +class SwishImplementation(torch.autograd.Function): + @staticmethod + def forward(ctx, i): + result = i * torch.sigmoid(i) + ctx.save_for_backward(i) + return result + + @staticmethod + def backward(ctx, grad_output): + i = ctx.saved_tensors[0] + sigmoid_i = torch.sigmoid(i) + return grad_output * (sigmoid_i * (1 + i * (1 - sigmoid_i))) + + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return SwishImplementation.apply(x) + + +def round_filters(filters, global_params): + """Calculate and round number of filters based on width multiplier. + Use width_coefficient, depth_divisor and min_depth of global_params. + + Args: + filters (int): Filters number to be calculated. + global_params (namedtuple): Global params of the model. + + Returns: + new_filters: New filters number after calculating. + """ + multiplier = global_params.width_coefficient + if not multiplier: + return filters + # TODO: modify the params names. + # maybe the names (width_divisor,min_width) + # are more suitable than (depth_divisor,min_depth). + divisor = global_params.depth_divisor + min_depth = global_params.min_depth + filters *= multiplier + min_depth = min_depth or divisor # pay attention to this line when using min_depth + # follow the formula transferred from official TensorFlow implementation + new_filters = max(min_depth, int(filters + divisor / 2) // divisor * divisor) + if new_filters < 0.9 * filters: # prevent rounding by more than 10% + new_filters += divisor + return int(new_filters) + + +def round_repeats(repeats, global_params): + """Calculate module's repeat number of a block based on depth multiplier. + Use depth_coefficient of global_params. + + Args: + repeats (int): num_repeat to be calculated. + global_params (namedtuple): Global params of the model. + + Returns: + new repeat: New repeat number after calculating. + """ + multiplier = global_params.depth_coefficient + if not multiplier: + return repeats + # follow the formula transferred from official TensorFlow implementation + return int(math.ceil(multiplier * repeats)) + + +def drop_connect(inputs, p, training): + """Drop connect. + + Args: + input (tensor: BCWH): Input of this structure. + p (float: 0.0~1.0): Probability of drop connection. + training (bool): The running mode. + + Returns: + output: Output after drop connection. + """ + assert 0 <= p <= 1, 'p must be in range of [0,1]' + + if not training: + return inputs + + batch_size = inputs.shape[0] + keep_prob = 1 - p + + # generate binary_tensor mask according to probability (p for 0, 1-p for 1) + random_tensor = keep_prob + random_tensor += torch.rand([batch_size, 1, 1, 1], dtype=inputs.dtype, device=inputs.device) + binary_tensor = torch.floor(random_tensor) + + output = inputs / keep_prob * binary_tensor + return output + + +def get_width_and_height_from_size(x): + """Obtain height and width from x. + + Args: + x (int, tuple or list): Data size. + + Returns: + size: A tuple or list (H,W). + """ + if isinstance(x, int): + return x, x + if isinstance(x, list) or isinstance(x, tuple): + return x + else: + raise TypeError() + + +def calculate_output_image_size(input_image_size, stride): + """Calculates the output image size when using Conv2dSamePadding with a stride. + Necessary for static padding. Thanks to mannatsingh for pointing this out. + + Args: + input_image_size (int, tuple or list): Size of input image. + stride (int, tuple or list): Conv2d operation's stride. + + Returns: + output_image_size: A list [H,W]. + """ + if input_image_size is None: + return None + image_height, image_width = get_width_and_height_from_size(input_image_size) + stride = stride if isinstance(stride, int) else stride[0] + image_height = int(math.ceil(image_height / stride)) + image_width = int(math.ceil(image_width / stride)) + return [image_height, image_width] + + +# Note: +# The following 'SamePadding' functions make output size equal ceil(input size/stride). +# Only when stride equals 1, can the output size be the same as input size. +# Don't be confused by their function names ! ! ! + +def get_same_padding_conv2d(image_size=None): + """Chooses static padding if you have specified an image size, and dynamic padding otherwise. + Static padding is necessary for ONNX exporting of models. + + Args: + image_size (int or tuple): Size of the image. + + Returns: + Conv2dDynamicSamePadding or Conv2dStaticSamePadding. + """ + if image_size is None: + return Conv2dDynamicSamePadding + else: + return partial(Conv2dStaticSamePadding, image_size=image_size) + + +class Conv2dDynamicSamePadding(nn.Conv2d): + """2D Convolutions like TensorFlow, for a dynamic image size. + The padding is operated in forward function by calculating dynamically. + """ + + # Tips for 'SAME' mode padding. + # Given the following: + # i: width or height + # s: stride + # k: kernel size + # d: dilation + # p: padding + # Output after Conv2d: + # o = floor((i+p-((k-1)*d+1))/s+1) + # If o equals i, i = floor((i+p-((k-1)*d+1))/s+1), + # => p = (i-1)*s+((k-1)*d+1)-i + + def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1, bias=True): + super().__init__(in_channels, out_channels, kernel_size, stride, 0, dilation, groups, bias) + self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2 + + def forward(self, x): + ih, iw = x.size()[-2:] + kh, kw = self.weight.size()[-2:] + sh, sw = self.stride + oh, ow = math.ceil(ih / sh), math.ceil(iw / sw) # change the output size according to stride ! ! ! + pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0) + pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0) + if pad_h > 0 or pad_w > 0: + x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]) + return F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups) + + +class Conv2dStaticSamePadding(nn.Conv2d): + """2D Convolutions like TensorFlow's 'SAME' mode, with the given input image size. + The padding mudule is calculated in construction function, then used in forward. + """ + + # With the same calculation as Conv2dDynamicSamePadding + + def __init__(self, in_channels, out_channels, kernel_size, stride=1, image_size=None, **kwargs): + super().__init__(in_channels, out_channels, kernel_size, stride, **kwargs) + self.stride = self.stride if len(self.stride) == 2 else [self.stride[0]] * 2 + + # Calculate padding based on image size and save it + assert image_size is not None + ih, iw = (image_size, image_size) if isinstance(image_size, int) else image_size + kh, kw = self.weight.size()[-2:] + sh, sw = self.stride + oh, ow = math.ceil(ih / sh), math.ceil(iw / sw) + pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0) + pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0) + if pad_h > 0 or pad_w > 0: + self.static_padding = nn.ZeroPad2d((pad_w // 2, pad_w - pad_w // 2, + pad_h // 2, pad_h - pad_h // 2)) + else: + self.static_padding = nn.Identity() + + def forward(self, x): + x = self.static_padding(x) + x = F.conv2d(x, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups) + return x + + +def get_same_padding_maxPool2d(image_size=None): + """Chooses static padding if you have specified an image size, and dynamic padding otherwise. + Static padding is necessary for ONNX exporting of models. + + Args: + image_size (int or tuple): Size of the image. + + Returns: + MaxPool2dDynamicSamePadding or MaxPool2dStaticSamePadding. + """ + if image_size is None: + return MaxPool2dDynamicSamePadding + else: + return partial(MaxPool2dStaticSamePadding, image_size=image_size) + + +class MaxPool2dDynamicSamePadding(nn.MaxPool2d): + """2D MaxPooling like TensorFlow's 'SAME' mode, with a dynamic image size. + The padding is operated in forward function by calculating dynamically. + """ + + def __init__(self, kernel_size, stride, padding=0, dilation=1, return_indices=False, ceil_mode=False): + super().__init__(kernel_size, stride, padding, dilation, return_indices, ceil_mode) + self.stride = [self.stride] * 2 if isinstance(self.stride, int) else self.stride + self.kernel_size = [self.kernel_size] * 2 if isinstance(self.kernel_size, int) else self.kernel_size + self.dilation = [self.dilation] * 2 if isinstance(self.dilation, int) else self.dilation + + def forward(self, x): + ih, iw = x.size()[-2:] + kh, kw = self.kernel_size + sh, sw = self.stride + oh, ow = math.ceil(ih / sh), math.ceil(iw / sw) + pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0) + pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0) + if pad_h > 0 or pad_w > 0: + x = F.pad(x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]) + return F.max_pool2d(x, self.kernel_size, self.stride, self.padding, + self.dilation, self.ceil_mode, self.return_indices) + + +class MaxPool2dStaticSamePadding(nn.MaxPool2d): + """2D MaxPooling like TensorFlow's 'SAME' mode, with the given input image size. + The padding mudule is calculated in construction function, then used in forward. + """ + + def __init__(self, kernel_size, stride, image_size=None, **kwargs): + super().__init__(kernel_size, stride, **kwargs) + self.stride = [self.stride] * 2 if isinstance(self.stride, int) else self.stride + self.kernel_size = [self.kernel_size] * 2 if isinstance(self.kernel_size, int) else self.kernel_size + self.dilation = [self.dilation] * 2 if isinstance(self.dilation, int) else self.dilation + + # Calculate padding based on image size and save it + assert image_size is not None + ih, iw = (image_size, image_size) if isinstance(image_size, int) else image_size + kh, kw = self.kernel_size + sh, sw = self.stride + oh, ow = math.ceil(ih / sh), math.ceil(iw / sw) + pad_h = max((oh - 1) * self.stride[0] + (kh - 1) * self.dilation[0] + 1 - ih, 0) + pad_w = max((ow - 1) * self.stride[1] + (kw - 1) * self.dilation[1] + 1 - iw, 0) + if pad_h > 0 or pad_w > 0: + self.static_padding = nn.ZeroPad2d((pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)) + else: + self.static_padding = nn.Identity() + + def forward(self, x): + x = self.static_padding(x) + x = F.max_pool2d(x, self.kernel_size, self.stride, self.padding, + self.dilation, self.ceil_mode, self.return_indices) + return x + + +################################################################################ +# Helper functions for loading model params +################################################################################ + +# BlockDecoder: A Class for encoding and decoding BlockArgs +# efficientnet_params: A function to query compound coefficient +# get_model_params and efficientnet: +# Functions to get BlockArgs and GlobalParams for efficientnet +# url_map and url_map_advprop: Dicts of url_map for pretrained weights +# load_pretrained_weights: A function to load pretrained weights + +class BlockDecoder(object): + """Block Decoder for readability, + straight from the official TensorFlow repository. + """ + + @staticmethod + def _decode_block_string(block_string): + """Get a block through a string notation of arguments. + + Args: + block_string (str): A string notation of arguments. + Examples: 'r1_k3_s11_e1_i32_o16_se0.25_noskip'. + + Returns: + BlockArgs: The namedtuple defined at the top of this file. + """ + assert isinstance(block_string, str) + + ops = block_string.split('_') + options = {} + for op in ops: + splits = re.split(r'(\d.*)', op) + if len(splits) >= 2: + key, value = splits[:2] + options[key] = value + + # Check stride + assert (('s' in options and len(options['s']) == 1) or + (len(options['s']) == 2 and options['s'][0] == options['s'][1])) + + return BlockArgs( + num_repeat=int(options['r']), + kernel_size=int(options['k']), + stride=[int(options['s'][0])], + expand_ratio=int(options['e']), + input_filters=int(options['i']), + output_filters=int(options['o']), + se_ratio=float(options['se']) if 'se' in options else None, + id_skip=('noskip' not in block_string)) + + @staticmethod + def _encode_block_string(block): + """Encode a block to a string. + + Args: + block (namedtuple): A BlockArgs type argument. + + Returns: + block_string: A String form of BlockArgs. + """ + args = [ + 'r%d' % block.num_repeat, + 'k%d' % block.kernel_size, + 's%d%d' % (block.strides[0], block.strides[1]), + 'e%s' % block.expand_ratio, + 'i%d' % block.input_filters, + 'o%d' % block.output_filters + ] + if 0 < block.se_ratio <= 1: + args.append('se%s' % block.se_ratio) + if block.id_skip is False: + args.append('noskip') + return '_'.join(args) + + @staticmethod + def decode(string_list): + """Decode a list of string notations to specify blocks inside the network. + + Args: + string_list (list[str]): A list of strings, each string is a notation of block. + + Returns: + blocks_args: A list of BlockArgs namedtuples of block args. + """ + assert isinstance(string_list, list) + blocks_args = [] + for block_string in string_list: + blocks_args.append(BlockDecoder._decode_block_string(block_string)) + return blocks_args + + @staticmethod + def encode(blocks_args): + """Encode a list of BlockArgs to a list of strings. + + Args: + blocks_args (list[namedtuples]): A list of BlockArgs namedtuples of block args. + + Returns: + block_strings: A list of strings, each string is a notation of block. + """ + block_strings = [] + for block in blocks_args: + block_strings.append(BlockDecoder._encode_block_string(block)) + return block_strings + + +def efficientnet_params(model_name): + """Map EfficientNet model name to parameter coefficients. + + Args: + model_name (str): Model name to be queried. + + Returns: + params_dict[model_name]: A (width,depth,res,dropout) tuple. + """ + params_dict = { + # Coefficients: width,depth,res,dropout + 'efficientnet-b0': (1.0, 1.0, 224, 0.2), + 'efficientnet-b1': (1.0, 1.1, 240, 0.2), + 'efficientnet-b2': (1.1, 1.2, 260, 0.3), + 'efficientnet-b3': (1.2, 1.4, 300, 0.3), + 'efficientnet-b4': (1.4, 1.8, 380, 0.4), + 'efficientnet-b5': (1.6, 2.2, 456, 0.4), + 'efficientnet-b6': (1.8, 2.6, 528, 0.5), + 'efficientnet-b7': (2.0, 3.1, 600, 0.5), + 'efficientnet-b8': (2.2, 3.6, 672, 0.5), + 'efficientnet-l2': (4.3, 5.3, 800, 0.5), + } + return params_dict[model_name] + + +def efficientnet(width_coefficient=None, depth_coefficient=None, image_size=None, + dropout_rate=0.2, drop_connect_rate=0.2, num_classes=1000, include_top=True): + """Create BlockArgs and GlobalParams for efficientnet model. + + Args: + width_coefficient (float) + depth_coefficient (float) + image_size (int) + dropout_rate (float) + drop_connect_rate (float) + num_classes (int) + + Meaning as the name suggests. + + Returns: + blocks_args, global_params. + """ + + # Blocks args for the whole model(efficientnet-b0 by default) + # It will be modified in the construction of EfficientNet Class according to model + blocks_args = [ + 'r1_k3_s11_e1_i32_o16_se0.25', + 'r2_k3_s22_e6_i16_o24_se0.25', + 'r2_k5_s22_e6_i24_o40_se0.25', + 'r3_k3_s22_e6_i40_o80_se0.25', + 'r3_k5_s11_e6_i80_o112_se0.25', + 'r4_k5_s22_e6_i112_o192_se0.25', + 'r1_k3_s11_e6_i192_o320_se0.25', + ] + blocks_args = BlockDecoder.decode(blocks_args) + + global_params = GlobalParams( + width_coefficient=width_coefficient, + depth_coefficient=depth_coefficient, + image_size=image_size, + dropout_rate=dropout_rate, + + num_classes=num_classes, + batch_norm_momentum=0.99, + batch_norm_epsilon=1e-3, + drop_connect_rate=drop_connect_rate, + depth_divisor=8, + min_depth=None, + include_top=include_top, + ) + + return blocks_args, global_params + + +def get_model_params(model_name, override_params): + """Get the block args and global params for a given model name. + + Args: + model_name (str): Model's name. + override_params (dict): A dict to modify global_params. + + Returns: + blocks_args, global_params + """ + if model_name.startswith('efficientnet'): + w, d, s, p = efficientnet_params(model_name) + # note: all models have drop connect rate = 0.2 + blocks_args, global_params = efficientnet( + width_coefficient=w, depth_coefficient=d, dropout_rate=p, image_size=s) + else: + raise NotImplementedError('model name is not pre-defined: {}'.format(model_name)) + if override_params: + # ValueError will be raised here if override_params has fields not included in global_params. + global_params = global_params._replace(**override_params) + return blocks_args, global_params + + +# train with Standard methods +# check more details in paper(EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks) +url_map = { + 'efficientnet-b0': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth', + 'efficientnet-b1': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b1-f1951068.pth', + 'efficientnet-b2': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b2-8bb594d6.pth', + 'efficientnet-b3': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b3-5fb5a3c3.pth', + 'efficientnet-b4': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b4-6ed6700e.pth', + 'efficientnet-b5': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b5-b6417697.pth', + 'efficientnet-b6': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b6-c76e70fd.pth', + 'efficientnet-b7': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b7-dcc49843.pth', +} + +# train with Adversarial Examples(AdvProp) +# check more details in paper(Adversarial Examples Improve Image Recognition) +url_map_advprop = { + 'efficientnet-b0': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b0-b64d5a18.pth', + 'efficientnet-b1': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b1-0f3ce85a.pth', + 'efficientnet-b2': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b2-6e9d97e5.pth', + 'efficientnet-b3': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b3-cdd7c0f4.pth', + 'efficientnet-b4': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b4-44fb3a87.pth', + 'efficientnet-b5': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b5-86493f6b.pth', + 'efficientnet-b6': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b6-ac80338e.pth', + 'efficientnet-b7': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b7-4652b6dd.pth', + 'efficientnet-b8': 'https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b8-22a8fe65.pth', +} + +# TODO: add the petrained weights url map of 'efficientnet-l2' + + +def load_pretrained_weights(model, model_name, weights_path=None, load_fc=True, advprop=False, verbose=True): + """Loads pretrained weights from weights path or download using url. + + Args: + model (Module): The whole model of efficientnet. + model_name (str): Model name of efficientnet. + weights_path (None or str): + str: path to pretrained weights file on the local disk. + None: use pretrained weights downloaded from the Internet. + load_fc (bool): Whether to load pretrained weights for fc layer at the end of the model. + advprop (bool): Whether to load pretrained weights + trained with advprop (valid when weights_path is None). + """ + if isinstance(weights_path, str): + state_dict = torch.load(weights_path) + else: + # AutoAugment or Advprop (different preprocessing) + url_map_ = url_map_advprop if advprop else url_map + state_dict = model_zoo.load_url(url_map_[model_name]) + + if load_fc: + ret = model.load_state_dict(state_dict, strict=False) + assert not ret.missing_keys, 'Missing keys when loading pretrained weights: {}'.format(ret.missing_keys) + else: + state_dict.pop('_fc.weight') + state_dict.pop('_fc.bias') + ret = model.load_state_dict(state_dict, strict=False) + assert set(ret.missing_keys) == set( + ['_fc.weight', '_fc.bias']), 'Missing keys when loading pretrained weights: {}'.format(ret.missing_keys) + assert not ret.unexpected_keys, 'Missing keys when loading pretrained weights: {}'.format(ret.unexpected_keys) + + if verbose: + print('Loaded pretrained weights for {}'.format(model_name)) diff --git a/ctlearn/core/pytorch/nets/models/NoPropDT/NoPropDT.py b/ctlearn/core/pytorch/nets/models/NoPropDT/NoPropDT.py new file mode 100644 index 00000000..113a9e46 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDT/NoPropDT.py @@ -0,0 +1,364 @@ +""" +NoProp-DT (No-Propagation Denoising Transformer) Model Module + +This module implements the NoProp-DT architecture, a denoising diffusion model +for classification tasks on Cherenkov telescope data. The model uses progressive +denoising through multiple diffusion steps to refine predictions. + +The NoProp-DT approach treats classification as a denoising task where noisy +class embeddings are progressively cleaned through T diffusion steps, ultimately +producing a clean classification prediction. + +Classes: + SimplifiedDenoiseBlock: Simplified denoise block for debugging/testing + NoPropDT: Main NoProp-DT model for classification with diffusion + +References: + - "Denoising Diffusion Probabilistic Models" (Ho et al., NeurIPS 2020) + - "Classifier-Free Diffusion Guidance" (Ho & Salimans, 2022) +""" + +import torch +from torch import nn +from .denoiseBlockThinRestNet import DenoiseBlock +import math + +class SimplifiedDenoiseBlock(nn.Module): + """ + Simplified denoising block for testing and debugging. + + This is a lightweight version of the full DenoiseBlock, useful for rapid + prototyping and debugging the diffusion process. It uses simple CNNs for + feature extraction and MLPs for processing embeddings. + + Architecture: + Image → CNN → Features (64-dim) + ↘ + Embedding → MLP → Features (256-dim) → Concat → MLP → Logits → Updated Embedding + + Attributes: + conv_path (nn.Sequential): CNN for image feature extraction + fc_z (nn.Sequential): MLP for embedding processing + combined (nn.Sequential): MLP for combining features and producing logits + """ + + def __init__(self, embedding_dim, num_classes): + """ + Initialize the simplified denoise block. + + Args: + embedding_dim (int): Dimension of class embeddings + num_classes (int): Number of output classes + """ + super().__init__() + # Simplified image feature extractor using CNN + self.conv_path = nn.Sequential( + nn.Conv2d(1, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), # Reduce spatial dimensions by 2x + nn.Conv2d(32, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), # Reduce spatial dimensions by 2x again + nn.AdaptiveAvgPool2d((1, 1)), # Global pooling to fixed size + nn.Flatten() # Flatten to (batch_size, 64) + ) + + # Simplified embedding processor using MLP + self.fc_z = nn.Sequential( + nn.Linear(embedding_dim, 256), + nn.ReLU() + ) + + # Combined processor: fuses image and embedding features + self.combined = nn.Sequential( + nn.Linear(256 + 64, 128), # Input: concatenated features + nn.ReLU(), + nn.Linear(128, num_classes) # Output: class logits + ) + + def forward(self, x, z_prev, W_embed): + """ + Forward pass through the simplified denoise block. + + Args: + x (torch.Tensor): Input images with shape (batch_size, 1, H, W) + z_prev (torch.Tensor): Previous embedding estimate (batch_size, embedding_dim) + W_embed (torch.Tensor): Class embedding matrix (num_classes, embedding_dim) + + Returns: + tuple: (z_next, logits) + - z_next: Updated embedding (batch_size, embedding_dim) + - logits: Class predictions (batch_size, num_classes) + """ + # Extract image features + x_feat = self.conv_path(x) # (batch_size, 64) + + # Process current embedding + z_feat = self.fc_z(z_prev) # (batch_size, 256) + + # Combine image and embedding features + combined = torch.cat([x_feat, z_feat], dim=1) # (batch_size, 320) + logits = self.combined(combined) # (batch_size, num_classes) + + # Update embedding using logits and class embeddings + # This pulls z_prev toward the class embedding indicated by logits + z_next = z_prev + logits @ W_embed # (batch_size, embedding_dim) + + return z_next, logits + +class NoPropDT(nn.Module): + """ + NoProp-DT: No-Propagation Denoising Transformer for classification. + + This model implements a diffusion-based approach to classification where + predictions are refined through multiple denoising steps. Each step removes + noise from class embeddings, progressively clarifying the prediction. + + Key Concepts: + + Diffusion Process: + - Forward: Add noise to true class embedding + - Reverse: Learn to denoise and recover true class + - Training: Match denoised output to clean embedding + - Inference: Start from noise, denoise T steps, classify + + Class Embeddings (W_embed): + - Each class has a learnable embedding vector + - These vectors represent "ideal" class representations + - Denoising process pulls noisy vectors toward these ideals + + Architecture Flow: + Noisy Embedding → [DenoiseBlock_1 → ... → DenoiseBlock_T] → Classifier → Prediction + + Attributes: + num_classes (int): Number of output classes + embedding_dim (int): Dimension of class embedding space + T (int): Number of diffusion steps + eta (float): Learning rate scaling factor for diffusion loss + W_embed (nn.Parameter): Learnable class embeddings (num_classes, embedding_dim) + blocks (nn.ModuleList): List of T denoising blocks + classifier (nn.Linear): Final classification head + alpha_bar (torch.Tensor): Noise schedule parameters + snr_diff (torch.Tensor): SNR differences for loss weighting + """ + + def __init__(self, num_outputs, embedding_dim=128, T=3, eta=0.1): + """ + Initialize the NoProp-DT model. + + Args: + num_outputs (int): Number of output classes + For telescope data: 2 (gamma vs proton) + embedding_dim (int, optional): Dimension of class embeddings. Defaults to 128 + Higher values allow richer representations but increase computation + T (int, optional): Number of diffusion steps. Defaults to 3 + More steps allow finer denoising but slower inference + Typical values: 3-10 + eta (float, optional): Diffusion loss scaling factor. Defaults to 0.1 + Controls the weight of denoising loss vs classification loss + + Example: + >>> # Create model for binary classification + >>> model = NoPropDT(num_outputs=2, embedding_dim=128, T=5) + >>> x = torch.randn(32, 1, 120, 120) # Batch of images + >>> output, _, _ = model(x) # Classification logits + >>> print(output.shape) # torch.Size([32, 2]) + """ + super().__init__() + + num_classes = num_outputs + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.T = T + self.eta = eta + + # Initialize learnable class embeddings + # Each class gets a random embedding vector that will be learned during training + # Small initial values (0.02 std) for stable training + self.W_embed = nn.Parameter( + torch.randn(num_classes, embedding_dim) * 0.02, + requires_grad=True + ) + + # Create T denoising blocks for progressive refinement + # Each block learns to remove one layer of noise + self.blocks = nn.ModuleList([ + DenoiseBlock(embedding_dim, num_classes) for _ in range(T) + ]) + + # Final classifier: maps clean embedding to class logits + self.classifier = nn.Linear(embedding_dim, num_classes) + + # Noise schedule: determines how much noise at each step + # Uses cosine schedule for smooth noise progression + self.register_buffer('alpha_bar', self._cosine_schedule(T)) + + # SNR differences: used to weight denoising loss at each step + # Steps with larger SNR improvements get higher weight + self.register_buffer('snr_diff', self._calculate_snr_diff(self.alpha_bar)) + + def _cosine_schedule(self, T): + """ + Generate a cosine noise schedule for diffusion. + + The cosine schedule provides smooth noise progression from high to low, + which has been shown to improve training stability compared to linear schedules. + + Mathematical Formula: + α̅_t = cos²((t/T + s) / (1 + s) × π/2) + + where s = 0.008 is a small offset to prevent numerical issues + + Args: + T (int): Total number of diffusion steps + + Returns: + torch.Tensor: Alpha bar values for each step with shape (T,) + Values decrease from ~1.0 (low noise) to ~0.0 (high noise) + + Properties: + - Monotonically decreasing: α̅_1 > α̅_2 > ... > α̅_T + - Smooth transitions between steps + - Prevents extreme noise levels at boundaries + """ + # Create timestep array from 1 to T + t = torch.arange(1, T + 1, dtype=torch.float32) + + # Apply cosine schedule with small offset for numerical stability + alpha_bar = torch.cos((t / T + 0.008) / 1.008 * (math.pi / 2)) ** 2 + + return alpha_bar + + def _calculate_snr_diff(self, alpha_bar): + """ + Calculate Signal-to-Noise Ratio differences for loss weighting. + + SNR differences quantify how much signal quality improves from one + diffusion step to the next. Steps with larger improvements receive + higher weight in the loss function. + + Mathematical Formula: + SNR_t = α̅_t / (1 - α̅_t) + SNR_diff_t = SNR_t - SNR_{t-1} + + Args: + alpha_bar (torch.Tensor): Noise schedule parameters with shape (T,) + + Returns: + torch.Tensor: SNR differences with shape (T,) + All values are positive (clamped to minimum 1e-5) + + Usage: + Used in training loss to weight each denoising step: + loss_t = snr_diff_t × MSE(denoised_t, target) + """ + # Calculate SNR at each step + snr = alpha_bar / (1 - alpha_bar + 1e-8) # Add epsilon for numerical stability + + # Prepend SNR_0 = 0 (no signal before first step) + snr_prev = torch.cat([torch.tensor([0.]), snr[:-1]]) + + # Compute differences and clamp to ensure positivity + return torch.clamp(snr - snr_prev, min=1e-5) + + def forward_denoise(self, x, z_prev, t): + """ + Perform one step of denoising at timestep t. + + This method applies the t-th denoising block to refine the current + embedding estimate. The block uses both the input image and the + current noisy embedding to produce a cleaner estimate. + + Args: + x (torch.Tensor): Input images with shape (batch_size, 1, H, W) + z_prev (torch.Tensor): Current noisy embedding (batch_size, embedding_dim) + t (int): Current timestep index (0 to T-1) + + Returns: + torch.Tensor: Refined embedding after denoising + Shape: (batch_size, embedding_dim) + Has less noise than z_prev + + Process: + 1. Denoise block extracts features from image + 2. Combines image features with current embedding + 3. Outputs refined embedding with reduced noise + """ + # Apply t-th denoising block + # Returns (denoised_embedding, logits), we only need the embedding + return self.blocks[t](x, z_prev, self.W_embed)[0] + + def inference(self, x): + """ + Perform full inference through all diffusion steps. + + This method executes the complete denoising process: + 1. Start from zero embedding (or random noise during training) + 2. Progressively denoise through T steps + 3. Classify the final clean embedding + + Args: + x (torch.Tensor): Input images with shape (batch_size, 1, H, W) + + Returns: + torch.Tensor: Classification logits with shape (batch_size, num_classes) + + Behavior: + Training (self.training=True): + - Can start from random noise for exploration + + Evaluation (self.training=False): + - Starts from zeros for deterministic predictions + + Example: + >>> model.eval() + >>> x = torch.randn(16, 1, 120, 120) + >>> logits = model.inference(x) + >>> predictions = torch.softmax(logits, dim=1) + >>> classes = predictions.argmax(dim=1) + """ + B = x.size(0) # Batch size + + # Initialize embedding + # Evaluation: Start from zeros for deterministic behavior + # Training: Could use noise for stochastic exploration (commented out) + z = torch.zeros(B, self.embedding_dim, device=x.device) + + # Progressive denoising through T steps + for t in range(self.T): + z = self.forward_denoise(x, z, t) + + # Classify the final clean embedding + return self.classifier(z) + + def forward(self, x): + """ + Forward pass through the model. + + This method provides a standardized interface compatible with other + models in CTLearn. It wraps the inference method and returns outputs + in the expected tuple format. + + Args: + x (torch.Tensor): Input images with shape (batch_size, 1, H, W) + + Returns: + tuple: (classification, energy, direction) where: + - classification: Logits for particle type (batch_size, num_classes) + - energy: None (not used for classification) + - direction: None (not used for classification) + + Note: + The tuple format (classification, energy, direction) is maintained + for compatibility with multi-task training frameworks, even though + this model only performs classification. + + Example: + >>> model = NoPropDT(num_outputs=2) + >>> x = torch.randn(32, 1, 120, 120) + >>> cls, energy, direction = model(x) + >>> print(cls.shape) # torch.Size([32, 2]) + >>> print(energy) # None + >>> print(direction) # None + """ + return self.inference(x), None, None \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/NoPropDT/denoiseBlockThinRestNet.py b/ctlearn/core/pytorch/nets/models/NoPropDT/denoiseBlockThinRestNet.py new file mode 100644 index 00000000..0915cbc3 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDT/denoiseBlockThinRestNet.py @@ -0,0 +1,402 @@ +""" +Denoising Block with Thin ResNet Architecture Module + +This module implements a denoising block architecture combining residual networks +with denoising capabilities for diffusion-based models. It's specifically designed +for the NoProp-DT (No-Propagation Denoising Transformer) model in CTLearn. + +The architecture uses a lightweight ResNet for feature extraction from images, +combined with MLP layers for processing noisy embeddings, making it suitable +for progressive denoising in diffusion models. + +Classes: + AdaptiveBatchNorm2d: Adaptive batch normalization with learnable parameters + ResidualBlock: Residual block with batch normalization and skip connections + DenoiseBlock: Main denoising block combining image and embedding features + MemoryEfficientSwish: Memory-efficient Swish activation function + +References: + - "Deep Residual Learning for Image Recognition" (He et al., CVPR 2016) + - "Denoising Diffusion Probabilistic Models" (Ho et al., NeurIPS 2020) +""" + +import torch +from torch import nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + """ + Memory-efficient implementation of Swish activation function. + + Swish (also known as SiLU - Sigmoid Linear Unit) is defined as: + f(x) = x · sigmoid(x) + + This implementation avoids storing intermediate activations during + forward pass, reducing memory consumption during backpropagation. + + Properties: + - Smooth and continuously differentiable + - Non-monotonic (can decrease for negative inputs) + - Self-gated: combines input with its sigmoid + - Better gradient flow than ReLU in some cases + """ + + def forward(self, x): + """ + Apply Swish activation element-wise. + + Args: + x (torch.Tensor): Input tensor of any shape + + Returns: + torch.Tensor: Activated tensor with same shape as input + """ + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + """ + Adaptive Batch Normalization with learnable interpolation parameters. + + This layer combines the input with batch-normalized input using learnable + parameters 'a' and 'b'. This allows the network to learn how much to rely + on the original input versus the normalized version. + + Mathematical Formula: + output = a * x + b * BatchNorm(x) + + where a and b are learnable scalar parameters initialized to 1 and 0 + respectively, so initially: output = x (identity mapping). + + Attributes: + bn (nn.BatchNorm2d): Standard batch normalization layer + a (nn.Parameter): Learnable weight for original input (initialized to 1) + b (nn.Parameter): Learnable weight for normalized input (initialized to 0) + + Benefits: + - Allows network to choose between normalized and unnormalized features + - Can help with training stability in diffusion models + - Provides more flexibility than standard batch normalization + """ + + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + """ + Initialize Adaptive Batch Normalization. + + Args: + num_features (int): Number of channels in the input + eps (float, optional): Small value for numerical stability. Defaults to 1e-5 + momentum (float, optional): Momentum for running statistics. Defaults to 0.5 + Note: Higher momentum (0.5 vs typical 0.1) gives more weight to current batch + affine (bool, optional): Whether to learn affine parameters. Defaults to True + """ + super(AdaptiveBatchNorm2d, self).__init__() + + # Standard batch normalization + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + + # Learnable interpolation parameters + # a=1, b=0 initially makes this an identity mapping + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + """ + Apply adaptive batch normalization. + + Process: + 1. Apply standard batch normalization to input + 2. Combine original input and normalized input using learnable weights + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, channels, height, width) + + Returns: + torch.Tensor: Adaptively normalized tensor with same shape as input + output = a * x + b * BatchNorm(x) + """ + return self.a * x + self.b * self.bn(x) + +class ResidualBlock(nn.Module): + """ + Residual block with adaptive batch normalization and skip connections. + + This block implements the core building block of ResNet architectures, + with two convolutional layers, adaptive batch normalization, and a + residual skip connection. If input and output channels differ, a 1x1 + convolution adjusts the skip connection. + + Architecture: + Input → [Conv3x3 → AdaptiveBN → ReLU → Conv3x3 → AdaptiveBN] → (+) → ReLU → Output + ↑ + | + Skip Connection + (1x1 conv if needed) + + Attributes: + conv_block (nn.Sequential): Main convolutional path with two conv layers + shortcut (nn.Sequential): Skip connection (identity or 1x1 conv) + relu (nn.ReLU): Final activation function + """ + + def __init__(self, in_channels, out_channels): + """ + Initialize the residual block. + + Args: + in_channels (int): Number of input channels + out_channels (int): Number of output channels + """ + super().__init__() + + # Main convolutional path + self.conv_block = nn.Sequential( + # First convolution: maintains spatial dimensions with padding + nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), + AdaptiveBatchNorm2d(out_channels), + nn.ReLU(), + # Second convolution: refines features + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + AdaptiveBatchNorm2d(out_channels) + ) + + # Skip connection: adjust channels if needed + self.shortcut = nn.Sequential() + if in_channels != out_channels: + # Use 1x1 convolution to match channel dimensions + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1), + nn.BatchNorm2d(out_channels) + ) + + # Final activation after adding residual + self.relu = nn.ReLU() + + def forward(self, x): + """ + Forward pass through the residual block. + + Process: + 1. Pass input through main convolutional path + 2. Add skip connection (possibly adjusted with 1x1 conv) + 3. Apply final ReLU activation + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, in_channels, H, W) + + Returns: + torch.Tensor: Output tensor with shape (batch_size, out_channels, H, W) + + Benefits of Residual Connection: + - Helps gradient flow during backpropagation + - Allows learning identity mappings when beneficial + - Enables training of very deep networks + """ + # Main path + skip connection + activation + return self.relu(self.conv_block(x) + self.shortcut(x)) + + +class DenoiseBlock(nn.Module): + """ + Denoising block combining image features and noisy embeddings. + + This block is the core component of the NoProp-DT diffusion model. It processes + telescope images through a lightweight ResNet while simultaneously processing + noisy class embeddings through MLPs, then combines both to produce a denoised + embedding and class logits. + + Architecture Overview: + Image → ThinResNet → Image Features (256-dim) + ↘ + Noisy Embedding → MLP (with residual) → Embedding Features (256-dim) → Concat → MLP → Logits + ↓ + z_next = z_prev + logits @ W_embed + + The ThinResNet Path: + Input (1, H, W) → + ResBlock(1→32) → MaxPool → Dropout → + ResBlock(32→64) → MaxPool → Dropout → + ResBlock(64→128) → AdaptiveAvgPool → + Flatten → Linear(128→256) + + The Embedding Path (with residual): + z_prev (embedding_dim) → + Linear(→256) → BN → ReLU → h1 → + Linear(→256) → BN → ReLU → h2 → + Linear(→256) → BN → h3 → + z_feat = h3 + h1 (residual) + + The Fusion Path: + Concat(image_feat, z_feat) → + Linear(512→256) → BN → ReLU → + Linear(256→128) → BN → ReLU → + Linear(128→num_classes) → logits + + Attributes: + use_softmax (bool): Whether to apply softmax to logits + conv_path (nn.Sequential): CNN for image feature extraction + fc_z1, fc_z2, fc_z3 (nn.Linear): MLP layers for embedding processing + bn_z1, bn_z2, bn_z3 (nn.BatchNorm1d): Batch norms for embedding MLP + fc_f1, fc_f2 (nn.Linear): MLP layers for fusing features + bn_f1, bn_f2 (nn.BatchNorm1d): Batch norms for fusion MLP + fc_out (nn.Linear): Output layer producing class logits + """ + + def __init__(self, embedding_dim, num_classes, use_softmax=False, num_channels=1): + """ + Initialize the denoising block. + + Args: + embedding_dim (int): Dimension of class embeddings + Typical values: 128, 256 + num_classes (int): Number of output classes + For telescope data: 2 (gamma vs proton) + use_softmax (bool, optional): Whether to apply softmax to logits. + Defaults to False. Set to True for probability outputs + num_channels (int, optional): Number of input image channels. + Defaults to 1 (grayscale telescope images) + """ + super().__init__() + self.use_softmax = use_softmax + + # ThinResNet convolutional path for image feature extraction + # Progressively increases channels: 1 → 32 → 64 → 128 + # Reduces spatial dimensions: H,W → H/2,W/2 → H/4,W/4 → 1,1 + self.conv_path = nn.Sequential( + # Stage 1: Initial feature extraction + ResidualBlock(num_channels, 32), + nn.MaxPool2d(2), # Downsample by 2x + nn.Dropout(0.2), # Regularization + + # Stage 2: Mid-level features + ResidualBlock(32, 64), + nn.MaxPool2d(2), # Downsample by 2x + nn.Dropout(0.2), + + # Stage 3: High-level features + ResidualBlock(64, 128), + + # Global pooling and projection + nn.AdaptiveAvgPool2d((1, 1)), # Reduce to (batch, 128, 1, 1) + nn.Flatten(), # (batch, 128) + nn.Linear(128, 256), # Project to 256-dim + ) + + # MLP for processing noisy embedding with residual connection + # Three-layer network with skip connection from first to last layer + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + # MLP for combining image and embedding features + # Fuses 512-dim (256 + 256) down to num_classes + self.fc_f1 = nn.Linear(256 + 256, 256) # Concatenated features + self.bn_f1 = nn.BatchNorm1d(256) + + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + + self.fc_out = nn.Linear(128, num_classes) # Final logits + + def forward(self, x, z_prev, W_embed): + """ + Forward pass through the denoising block. + + This method combines image features with noisy embeddings to produce + both denoised embeddings and class predictions. The denoising happens + through the interaction between image features and the noisy embedding. + + Process: + 1. Extract features from image using ThinResNet + 2. Process noisy embedding through MLP with residual connection + 3. Concatenate image and embedding features + 4. Generate class logits through fusion MLP + 5. Update embedding: z_next = z_prev + logits @ W_embed + + Args: + x (torch.Tensor): Input images + Shape: (batch_size, num_channels, height, width) + Typically: (batch_size, 1, 120, 120) for telescope images + + z_prev (torch.Tensor): Previous/current noisy embedding + Shape: (batch_size, embedding_dim) + Contains noise that will be reduced in this step + + W_embed (torch.Tensor): Class embedding matrix + Shape: (num_classes, embedding_dim) + Each row is the embedding for one class + + Returns: + tuple: (z_next, logits) where: + - z_next (torch.Tensor): Denoised embedding + Shape: (batch_size, embedding_dim) + Cleaner than z_prev, closer to true class embedding + - logits (torch.Tensor): Class predictions + Shape: (batch_size, num_classes) + If use_softmax=True: softmax probabilities + If use_softmax=False: raw logits + + Denoising Mechanism: + The update rule z_next = z_prev + logits @ W_embed works by: + 1. logits indicate which class is most likely + 2. W_embed provides the ideal embedding for each class + 3. The product logits @ W_embed pulls z_prev toward the correct class embedding + 4. Over multiple diffusion steps, this progressively cleans the embedding + + Example: + >>> block = DenoiseBlock(embedding_dim=128, num_classes=2) + >>> x = torch.randn(32, 1, 120, 120) # Batch of 32 images + >>> z_prev = torch.randn(32, 128) # Noisy embeddings + >>> W_embed = torch.randn(2, 128) # Class embeddings + >>> z_next, logits = block(x, z_prev, W_embed) + >>> print(z_next.shape) # torch.Size([32, 128]) + >>> print(logits.shape) # torch.Size([32, 2]) + """ + # Step 1: Extract features from input image + # Shape: (batch_size, 256) + x_feat = self.conv_path(x) + + # Step 2: Process noisy embedding through MLP with residual connection + # First layer + h1 = F.relu(self.bn_z1(self.fc_z1(z_prev))) # (batch_size, 256) + + # Second layer + h2 = F.relu(self.bn_z2(self.fc_z2(h1))) # (batch_size, 256) + + # Third layer (no ReLU here) + h3 = self.bn_z3(self.fc_z3(h2)) # (batch_size, 256) + + # Residual connection: add h1 to h3 + # This helps gradient flow and preserves early features + z_feat = h3 + h1 # (batch_size, 256) + + # Step 3: Concatenate image and embedding features + # Combines information from both modalities + h_f = torch.cat([x_feat, z_feat], dim=1) # (batch_size, 512) + + # Step 4: Process combined features through fusion MLP + # First fusion layer + h_f = F.relu(self.bn_f1(self.fc_f1(h_f))) # (batch_size, 256) + + # Second fusion layer + h_f = F.relu(self.bn_f2(self.fc_f2(h_f))) # (batch_size, 128) + + # Step 5: Compute logits for all classes + logits = self.fc_out(h_f) # (batch_size, num_classes) + + # Step 6: Optionally convert logits to probabilities + if self.use_softmax: + p = F.softmax(logits, dim=1) + else: + p = logits + + # Step 7: Compute next denoised embedding + # Update rule: z_next = z_prev + logits @ W_embed + # This pulls z_prev toward the embedding of the predicted class + z_next = z_prev + logits @ W_embed # (batch_size, embedding_dim) + + return z_next, logits diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg/NoPropDTReg.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg/NoPropDTReg.py new file mode 100644 index 00000000..89ba33aa --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg/NoPropDTReg.py @@ -0,0 +1,121 @@ +# NoProp-DT model + +import torch +from torch import nn +import torch.nn.functional as F +# from .denoiseBlock import DenoiseBlock +from .denoiseBlockThinRestNet import DenoiseBlock, MemoryEfficientSwish + +import math + +class SimplifiedDenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_classes): + super().__init__() + # Simplified image feature extractor + self.conv_path = nn.Sequential( + nn.Conv2d(1, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.Conv2d(32, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten() + ) + + # Simplified embedding processor + self.fc_z = nn.Sequential( + nn.Linear(embedding_dim, 256), + nn.ReLU() + ) + + # Combined processor + self.combined = nn.Sequential( + nn.Linear(256 + 64, 128), + nn.ReLU(), + nn.Linear(128, num_classes) + ) + + def forward(self, x, z_prev, W_embed): + # Image features + x_feat = self.conv_path(x) + + # Process embedding + z_feat = self.fc_z(z_prev) + + # Combine features + combined = torch.cat([x_feat, z_feat], dim=1) + logits = self.combined(combined) + + # Update embedding + z_next = z_prev + logits @ W_embed + + return z_next, logits + +class NoPropDTReg(nn.Module): + def __init__(self, task, num_outputs, embedding_dim=128, T=3, eta=0.1,num_blocks=[2, 3, 3, 3]): + super().__init__() + + self.task = task + num_classes = num_outputs + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.T = T + self.eta = eta + + self.blocks = nn.ModuleList([DenoiseBlock(embedding_dim,num_channels=1,num_blocks=num_blocks) for _ in range(T)]) + + self.regressor = nn.Sequential( + nn.Linear(embedding_dim, embedding_dim//2), + # MemoryEfficientSwish(), + nn.Linear(embedding_dim//2, num_outputs) +) + + # Final classifier + self.classifier = nn.Linear(embedding_dim, num_classes) + + # Improved noise schedule + self.register_buffer('alpha_bar', self._cosine_schedule(T)) + self.register_buffer('snr_diff', self._calculate_snr_diff(self.alpha_bar)) + + + self.target_embedder = nn.Linear(num_outputs, embedding_dim) + + for m in self.regressor: + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + nn.init.zeros_(m.bias) + + def _cosine_schedule(self, T): + t = torch.arange(1, T+1, dtype=torch.float32) + alpha_bar = torch.cos((t / T + 0.008) / 1.008 * (math.pi/2))**2 + return alpha_bar + + def _calculate_snr_diff(self, alpha_bar): + snr = alpha_bar / (1 - alpha_bar + 1e-8) + snr_prev = torch.cat([torch.tensor([0.]), snr[:-1]]) + return torch.clamp(snr - snr_prev, min=1e-5) + + def forward_denoise(self, x, z_prev, t): + return self.blocks[t](x, z_prev, self.target_embedder)[0] + + def regress(self, z): + return self.regressor(z) + + def inference(self, x): + B = x.size(0) + z = torch.randn(B, self.embedding_dim, device=x.device) + if not self.training: + z = torch.zeros(B, self.embedding_dim, device=x.device) + + for t in range(self.T): + z = self.forward_denoise(x, z, t) + + return self.regress(z) + + def forward(self, x): + + if self.task=="direction": + return None, None, self.inference(x) + elif self.task=="energy": + return None, F.tanh(self.inference(x))*4, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet copy 2.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet copy 2.py new file mode 100644 index 00000000..f1326018 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet copy 2.py @@ -0,0 +1,116 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn1 = AdaptiveBatchNorm2d(out_channels) + self.act1 = nn.GELU() + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn2 = AdaptiveBatchNorm2d(out_channels) + self.act2 = nn.GELU() + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + def forward(self, x): + # residual = self.act1(self.bn1(self.conv1(x))) + # residual = self.act2(self.bn2(self.conv2(residual))) + residual = self.act1((self.conv1(x))) + residual = self.act2((self.conv2(residual))) + residual = self.se(residual) + out = residual + self.shortcut(x) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, drop_prob=0.2): + super().__init__() + # Ahora más profundo y ancho: + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(256, 512), # Embedding más grande + # nn.BatchNorm1d(512), + nn.GELU() + ) + + self.fc_z1 = nn.Linear(embedding_dim, 512) + self.bn_z1 = nn.BatchNorm1d(512) + self.fc_z2 = nn.Linear(512, 512) + self.bn_z2 = nn.BatchNorm1d(512) + self.fc_z3 = nn.Linear(512, 512) + self.bn_z3 = nn.BatchNorm1d(512) + self.fc_f1 = nn.Linear(1024, 512) + self.bn_f1 = nn.BatchNorm1d(512) + self.fc_f2 = nn.Linear(512, 256) + self.bn_f2 = nn.BatchNorm1d(256) + self.fc_out = nn.Linear(256, embedding_dim) + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + # h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + # h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h1 = self.act1((self.fc_z1(z_prev))) + h2 = self.act2((self.fc_z2(h1))) + + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + # h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + # h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + h_f = self.act_f1((self.fc_f1(h_f))) + h_f = self.act_f2((self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet copy.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet copy.py new file mode 100644 index 00000000..6487c4ee --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet copy.py @@ -0,0 +1,111 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn1 = AdaptiveBatchNorm2d(out_channels) + self.act1 = nn.GELU() + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn2 = AdaptiveBatchNorm2d(out_channels) + self.act2 = nn.GELU() + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + def forward(self, x): + # residual = self.act1(self.bn1(self.conv1(x))) + # residual = self.act2(self.bn2(self.conv2(residual))) + residual = self.act1((self.conv1(x))) + residual = self.act2((self.conv2(residual))) + residual = self.se(residual) + out = residual + self.shortcut(x) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, drop_prob=0.2): + super().__init__() + # Ahora más profundo y ancho: + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(256, 512), # Embedding más grande + nn.BatchNorm1d(512), + nn.GELU() + ) + + self.fc_z1 = nn.Linear(embedding_dim, 512) + self.bn_z1 = nn.BatchNorm1d(512) + self.fc_z2 = nn.Linear(512, 512) + self.bn_z2 = nn.BatchNorm1d(512) + self.fc_z3 = nn.Linear(512, 512) + self.bn_z3 = nn.BatchNorm1d(512) + self.fc_f1 = nn.Linear(1024, 512) + self.bn_f1 = nn.BatchNorm1d(512) + self.fc_f2 = nn.Linear(512, 256) + self.bn_f2 = nn.BatchNorm1d(256) + self.fc_out = nn.Linear(256, embedding_dim) + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet.py new file mode 100644 index 00000000..c9768565 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet.py @@ -0,0 +1,158 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channel, reduction=16): + super(SEBlock, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.avg_pool(x).squeeze(-1).squeeze(-1) # Ensuring dimension match + y = self.fc(y).view(b, c, 1, 1) + return x * y.expand_as(x) + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_channels, out_channels, stride=1, reduction=16,use_bn=False): + + + super(BasicBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) + self.use_bn = use_bn + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) + + + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != self.expansion * out_channels: + + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(self.expansion * out_channels) + ) + + def forward(self, x): + out = F.relu((self.conv1(x))) + out = (self.conv2(out)) + + # print("out.shape", out.shape, "shortcut.shape", self.shortcut(x).shape) + # print("out.dtype", out.dtype, "shortcut.dtype", self.shortcut(x).dtype) + # print("out.device", out.device, "shortcut.device", self.shortcut(x).device) + # print("out.is_contiguous()", out.is_contiguous(), "shortcut.is_contiguous()", self.shortcut(x).is_contiguous()) + + # assert out.shape == self.shortcut(x).shape, f"Shape mismatch: {out.shape} vs {self.shortcut(x).shape}" + # assert out.dtype == self.shortcut(x).dtype, f"Dtype mismatch: {out.dtype} vs {self.shortcut(x).dtype}" + # assert out.device == self.shortcut(x).device, f"Device mismatch: {out.device} vs {self.shortcut(x).device}" + + # out = out.contiguous() + + out += self.shortcut(x).contiguous() + out = F.relu(out) + out = self.se(out) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, block=BasicBlock, num_blocks=[2, 3, 3, 3], num_inputs=1, num_outputs=2,use_bn=False,dropout=0.0): + super().__init__() + + # block = BasicBlock + self.in_channels = 64 + self.use_bn=use_bn + + self.conv1 = nn.Conv2d(num_channels, 64, kernel_size=3, stride=1, padding=1, bias=False) + + self.layer1_1 = self._make_layer(block, 32, num_blocks[0], stride=1) + self.layer2_1 = self._make_layer(block, 64, num_blocks[1], stride=2) + self.layer3_1 = self._make_layer(block, 128, num_blocks[2], stride=2) + self.layer4_1 = self._make_layer(block, 256, num_blocks[3], stride=2) + # Reducing the number of layers and filters to make it "thin" + # self.fc_1 = nn.Linear(embedding_dim , embedding_dim) + # self.fc_2 = nn.Linear(embedding_dim , 256) + self.bn_final = nn.BatchNorm1d(512 * block.expansion) # BatchNorm layer + self.prelu = nn.PReLU(num_parameters=512 * block.expansion) # Define Leaky ReLU + + + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + self.dropout = nn.Dropout(dropout) + + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + self.fc_f1 = nn.Linear(512, 256) + self.bn_f1 = nn.BatchNorm1d(256) + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + self.fc_out = nn.Linear(128, embedding_dim) + + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + + def forward(self, x, z_prev, _): + + out_1 = F.relu(self.conv1(x)) + out_1 = self.layer1_1(out_1) + out_1 = self.layer2_1(out_1) + out_1 = self.layer3_1(out_1) + x_feat = self.layer4_1(out_1) + + x_feat= self.adaptive_pool(x_feat) + x_feat = x_feat.view(x_feat.size(0), -1) + + + # h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + # h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h1 = self.act1((self.fc_z1(z_prev))) + h2 = self.act2((self.fc_z2(h1))) + + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + # h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + # h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + h_f = self.act_f1((self.fc_f1(h_f))) + h_f = self.act_f2((self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None + + def _make_layer(self, block, out_channels, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_channels, out_channels, stride, use_bn=self.use_bn)) + self.in_channels = out_channels * block.expansion + return nn.Sequential(*layers) + diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet_original.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet_original.py new file mode 100644 index 00000000..fc52d51f --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg/denoiseBlockThinRestNet_original.py @@ -0,0 +1,154 @@ +# Denoising block +import torch +from torch import nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +# class ResidualBlock(nn.Module): +# def __init__(self, in_channels, out_channels): +# super().__init__() +# self.conv_block = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels), +# nn.PReLU(), +# nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels) +# ) + +# self.shortcut = nn.Sequential() +# if in_channels != out_channels: +# self.shortcut = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=1), +# nn.BatchNorm2d(out_channels) +# ) + +# self.relu = nn.PReLU() + +# def forward(self, x): +# return self.relu(self.conv_block(x) + self.shortcut(x)) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16, drop_path_rate=0.1): + super().__init__() + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.norm1 = nn.GroupNorm(8, out_channels) + self.act1 = nn.GELU() + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + self.norm2 = nn.GroupNorm(8, out_channels) + self.act2 = nn.GELU() + + self.se = SEBlock(out_channels, reduction) + + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1), + nn.GroupNorm(8, out_channels) + ) + + try: + from timm.models.layers import DropPath + self.drop_path = DropPath(drop_path_rate) + except ImportError: + self.drop_path = nn.Identity() + + # self.final_act = nn.GELU() + self.final_act = MemoryEfficientSwish() + def forward(self, x): + residual = self.act1(self.norm1(self.conv1(x))) + residual = self.act2(self.norm2(self.conv2(residual))) + # residual = self.act1((self.conv1(x))) + # residual = self.act2((self.conv2(residual))) + + residual = self.se(residual) + + out = self.drop_path(residual) + self.shortcut(x) + return self.final_act(out) + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1): + super().__init__() + + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 32), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(32, 64), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(64, 128), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(128, 256), + ) + + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + self.fc_f1 = nn.Linear(512, 256) + self.bn_f1 = nn.BatchNorm1d(256) + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + self.fc_out = nn.Linear(128, embedding_dim) + + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + + z_feat = h3 + h1 + + h_f = torch.cat([x_feat, z_feat], dim=1) + + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg2/NoPropDTReg2.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/NoPropDTReg2.py new file mode 100644 index 00000000..a5a7c8f9 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/NoPropDTReg2.py @@ -0,0 +1,121 @@ +# NoProp-DT model + +import torch +from torch import nn + +# from .denoiseBlock import DenoiseBlock +from .denoiseBlockThinRestNet import DenoiseBlock, MemoryEfficientSwish + +import math + +class SimplifiedDenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_classes): + super().__init__() + # Simplified image feature extractor + self.conv_path = nn.Sequential( + nn.Conv2d(1, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.Conv2d(32, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten() + ) + + # Simplified embedding processor + self.fc_z = nn.Sequential( + nn.Linear(embedding_dim, 256), + nn.ReLU() + ) + + # Combined processor + self.combined = nn.Sequential( + nn.Linear(256 + 64, 128), + nn.ReLU(), + nn.Linear(128, num_classes) + ) + + def forward(self, x, z_prev, W_embed): + # Image features + x_feat = self.conv_path(x) + + # Process embedding + z_feat = self.fc_z(z_prev) + + # Combine features + combined = torch.cat([x_feat, z_feat], dim=1) + logits = self.combined(combined) + + # Update embedding + z_next = z_prev + logits @ W_embed + + return z_next, logits + +class NoPropDTReg2(nn.Module): + def __init__(self, task, num_outputs, embedding_dim=128, T=3, eta=0.1,num_blocks=[2, 3, 3, 3]): + super().__init__() + + self.task = task + num_classes = num_outputs + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.T = T + self.eta = eta + + self.blocks = nn.ModuleList([DenoiseBlock(embedding_dim,num_channels=1,num_blocks=num_blocks) for _ in range(T)]) + + self.regressor = nn.Sequential( + nn.Linear(embedding_dim, embedding_dim//2), + # MemoryEfficientSwish(), + nn.Linear(embedding_dim//2, num_outputs) + ) + + # Final classifier + self.classifier = nn.Linear(embedding_dim, num_classes) + + # Improved noise schedule + self.register_buffer('alpha_bar', self._cosine_schedule(T)) + self.register_buffer('snr_diff', self._calculate_snr_diff(self.alpha_bar)) + + + self.target_embedder = nn.Linear(num_outputs, embedding_dim) + + for m in self.regressor: + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + nn.init.zeros_(m.bias) + + def _cosine_schedule(self, T): + t = torch.arange(1, T+1, dtype=torch.float32) + alpha_bar = torch.cos((t / T + 0.008) / 1.008 * (math.pi/2))**2 + return alpha_bar + + def _calculate_snr_diff(self, alpha_bar): + snr = alpha_bar / (1 - alpha_bar + 1e-8) + snr_prev = torch.cat([torch.tensor([0.]), snr[:-1]]) + return torch.clamp(snr - snr_prev, min=1e-5) + + def forward_denoise(self, x, z_prev, t): + return self.blocks[t](x, z_prev, None)[0] + + def regress(self, z): + return self.regressor(z) + + def inference(self, x): + B = x.size(0) + z = torch.randn(B, self.embedding_dim, device=x.device) + if not self.training: + z = torch.zeros(B, self.embedding_dim, device=x.device) + + for t in range(self.T): + z = self.forward_denoise(x, z, t) + + return self.regress(z) + + def forward(self, x): + + if self.task=="direction": + return None, None, self.inference(x) + elif self.task=="energy": + return None, self.inference(x), None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet copy 2.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet copy 2.py new file mode 100644 index 00000000..f1326018 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet copy 2.py @@ -0,0 +1,116 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn1 = AdaptiveBatchNorm2d(out_channels) + self.act1 = nn.GELU() + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn2 = AdaptiveBatchNorm2d(out_channels) + self.act2 = nn.GELU() + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + def forward(self, x): + # residual = self.act1(self.bn1(self.conv1(x))) + # residual = self.act2(self.bn2(self.conv2(residual))) + residual = self.act1((self.conv1(x))) + residual = self.act2((self.conv2(residual))) + residual = self.se(residual) + out = residual + self.shortcut(x) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, drop_prob=0.2): + super().__init__() + # Ahora más profundo y ancho: + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(256, 512), # Embedding más grande + # nn.BatchNorm1d(512), + nn.GELU() + ) + + self.fc_z1 = nn.Linear(embedding_dim, 512) + self.bn_z1 = nn.BatchNorm1d(512) + self.fc_z2 = nn.Linear(512, 512) + self.bn_z2 = nn.BatchNorm1d(512) + self.fc_z3 = nn.Linear(512, 512) + self.bn_z3 = nn.BatchNorm1d(512) + self.fc_f1 = nn.Linear(1024, 512) + self.bn_f1 = nn.BatchNorm1d(512) + self.fc_f2 = nn.Linear(512, 256) + self.bn_f2 = nn.BatchNorm1d(256) + self.fc_out = nn.Linear(256, embedding_dim) + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + # h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + # h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h1 = self.act1((self.fc_z1(z_prev))) + h2 = self.act2((self.fc_z2(h1))) + + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + # h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + # h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + h_f = self.act_f1((self.fc_f1(h_f))) + h_f = self.act_f2((self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet copy.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet copy.py new file mode 100644 index 00000000..6487c4ee --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet copy.py @@ -0,0 +1,111 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn1 = AdaptiveBatchNorm2d(out_channels) + self.act1 = nn.GELU() + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn2 = AdaptiveBatchNorm2d(out_channels) + self.act2 = nn.GELU() + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + def forward(self, x): + # residual = self.act1(self.bn1(self.conv1(x))) + # residual = self.act2(self.bn2(self.conv2(residual))) + residual = self.act1((self.conv1(x))) + residual = self.act2((self.conv2(residual))) + residual = self.se(residual) + out = residual + self.shortcut(x) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, drop_prob=0.2): + super().__init__() + # Ahora más profundo y ancho: + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(256, 512), # Embedding más grande + nn.BatchNorm1d(512), + nn.GELU() + ) + + self.fc_z1 = nn.Linear(embedding_dim, 512) + self.bn_z1 = nn.BatchNorm1d(512) + self.fc_z2 = nn.Linear(512, 512) + self.bn_z2 = nn.BatchNorm1d(512) + self.fc_z3 = nn.Linear(512, 512) + self.bn_z3 = nn.BatchNorm1d(512) + self.fc_f1 = nn.Linear(1024, 512) + self.bn_f1 = nn.BatchNorm1d(512) + self.fc_f2 = nn.Linear(512, 256) + self.bn_f2 = nn.BatchNorm1d(256) + self.fc_out = nn.Linear(256, embedding_dim) + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet.py new file mode 100644 index 00000000..a1cd048d --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet.py @@ -0,0 +1,146 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channel, reduction=16): + super(SEBlock, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.avg_pool(x).squeeze(-1).squeeze(-1) # Ensuring dimension match + y = self.fc(y).view(b, c, 1, 1) + return x * y.expand_as(x) + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_channels, out_channels, stride=1, reduction=16,use_bn=False): + + + super(BasicBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) + self.use_bn = use_bn + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) + + + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != self.expansion * out_channels: + + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(self.expansion * out_channels) + ) + + def forward(self, x): + out = F.relu((self.conv1(x))) + out = (self.conv2(out)) + out += self.shortcut(x) + out = F.relu(out) + out = self.se(out) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, block=BasicBlock, num_blocks=[2, 3, 3, 3], num_inputs=1, num_outputs=2,use_bn=False,dropout=0.0): + super().__init__() + + # block = BasicBlock + self.in_channels = 64 + self.use_bn=use_bn + + self.conv1 = nn.Conv2d(num_channels, 64, kernel_size=3, stride=1, padding=1, bias=False) + + self.layer1_1 = self._make_layer(block, 32, num_blocks[0], stride=1) + self.layer2_1 = self._make_layer(block, 64, num_blocks[1], stride=2) + self.layer3_1 = self._make_layer(block, 128, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block, 256, num_blocks[3], stride=2) + # Reducing the number of layers and filters to make it "thin" + # self.fc_1 = nn.Linear(embedding_dim , embedding_dim) + # self.fc_2 = nn.Linear(embedding_dim , 256) + self.bn_final = nn.BatchNorm1d(512 * block.expansion) # BatchNorm layer + self.prelu = nn.PReLU(num_parameters=512 * block.expansion) # Define Leaky ReLU + + + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + self.dropout = nn.Dropout(dropout) + + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + self.fc_f1 = nn.Linear(512, 256) + self.bn_f1 = nn.BatchNorm1d(256) + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + self.fc_out = nn.Linear(128, embedding_dim) + + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + + def forward(self, x, z_prev, _): + + out_1 = F.relu(self.conv1(x)) + out_1 = self.layer1_1(out_1) + out_1 = self.layer2_1(out_1) + out_1 = self.layer3_1(out_1) + + x_feat = self.layer4(out_1) + x_feat= self.adaptive_pool(x_feat) + x_feat = x_feat.view(x_feat.size(0), -1) + + + # h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + # h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h1 = self.act1((self.fc_z1(z_prev))) + h2 = self.act2((self.fc_z2(h1))) + + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + # h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + # h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + h_f = self.act_f1((self.fc_f1(h_f))) + h_f = self.act_f2((self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None + + def _make_layer(self, block, out_channels, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_channels, out_channels, stride, use_bn=self.use_bn)) + self.in_channels = out_channels * block.expansion + return nn.Sequential(*layers) + diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet_original.py b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet_original.py new file mode 100644 index 00000000..fc52d51f --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTReg2/denoiseBlockThinRestNet_original.py @@ -0,0 +1,154 @@ +# Denoising block +import torch +from torch import nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +# class ResidualBlock(nn.Module): +# def __init__(self, in_channels, out_channels): +# super().__init__() +# self.conv_block = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels), +# nn.PReLU(), +# nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels) +# ) + +# self.shortcut = nn.Sequential() +# if in_channels != out_channels: +# self.shortcut = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=1), +# nn.BatchNorm2d(out_channels) +# ) + +# self.relu = nn.PReLU() + +# def forward(self, x): +# return self.relu(self.conv_block(x) + self.shortcut(x)) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16, drop_path_rate=0.1): + super().__init__() + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.norm1 = nn.GroupNorm(8, out_channels) + self.act1 = nn.GELU() + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + self.norm2 = nn.GroupNorm(8, out_channels) + self.act2 = nn.GELU() + + self.se = SEBlock(out_channels, reduction) + + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1), + nn.GroupNorm(8, out_channels) + ) + + try: + from timm.models.layers import DropPath + self.drop_path = DropPath(drop_path_rate) + except ImportError: + self.drop_path = nn.Identity() + + # self.final_act = nn.GELU() + self.final_act = MemoryEfficientSwish() + def forward(self, x): + residual = self.act1(self.norm1(self.conv1(x))) + residual = self.act2(self.norm2(self.conv2(residual))) + # residual = self.act1((self.conv1(x))) + # residual = self.act2((self.conv2(residual))) + + residual = self.se(residual) + + out = self.drop_path(residual) + self.shortcut(x) + return self.final_act(out) + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1): + super().__init__() + + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 32), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(32, 64), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(64, 128), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(128, 256), + ) + + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + self.fc_f1 = nn.Linear(512, 256) + self.bn_f1 = nn.BatchNorm1d(256) + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + self.fc_out = nn.Linear(128, embedding_dim) + + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + + z_feat = h3 + h1 + + h_f = torch.cat([x_feat, z_feat], dim=1) + + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/NoPropDTRegDBB.py b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/NoPropDTRegDBB.py new file mode 100644 index 00000000..a9cf5623 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/NoPropDTRegDBB.py @@ -0,0 +1,125 @@ +# NoProp-DT model + +import torch +from torch import nn +import torch.nn.functional as F +# from .denoiseBlock import DenoiseBlock +from .denoiseBlockThinRestNet import DenoiseBlock, MemoryEfficientSwish + +import math + +class SimplifiedDenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_classes): + super().__init__() + # Simplified image feature extractor + self.conv_path = nn.Sequential( + nn.Conv2d(1, 32, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.Conv2d(32, 64, kernel_size=3, padding=1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten() + ) + + # Simplified embedding processor + self.fc_z = nn.Sequential( + nn.Linear(embedding_dim, 256), + nn.ReLU() + ) + + # Combined processor + self.combined = nn.Sequential( + nn.Linear(256 + 64, 128), + nn.ReLU(), + nn.Linear(128, num_classes) + ) + + def forward(self, x, z_prev, W_embed): + # Image features + x_feat = self.conv_path(x) + + # Process embedding + z_feat = self.fc_z(z_prev) + + # Combine features + combined = torch.cat([x_feat, z_feat], dim=1) + logits = self.combined(combined) + + # Update embedding + z_next = z_prev + logits @ W_embed + + return z_next, logits + +class NoPropDTRegDBB(nn.Module): + def __init__(self, task, num_outputs, embedding_dim=128, T=3, eta=0.1,num_blocks=[2, 3, 3, 3]): + super().__init__() + + self.task = task + num_classes = num_outputs + self.num_classes = num_classes + self.embedding_dim = embedding_dim + self.T = T + self.eta = eta + + self.blocks = nn.ModuleList([DenoiseBlock(embedding_dim,num_channels=1,num_blocks=num_blocks) for _ in range(T)]) + + self.regressor = nn.Sequential( + nn.Linear(embedding_dim, embedding_dim//2), + # MemoryEfficientSwish(), + nn.Linear(embedding_dim//2, num_outputs) +) + + # Final classifier + self.classifier = nn.Linear(embedding_dim, num_classes) + + # Improved noise schedule + self.register_buffer('alpha_bar', self._cosine_schedule(T)) + self.register_buffer('snr_diff', self._calculate_snr_diff(self.alpha_bar)) + + + self.target_embedder = nn.Linear(num_outputs, embedding_dim) + + for m in self.regressor: + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + nn.init.zeros_(m.bias) + + def _cosine_schedule(self, T): + t = torch.arange(1, T+1, dtype=torch.float32) + alpha_bar = torch.cos((t / T + 0.008) / 1.008 * (math.pi/2))**2 + return alpha_bar + + def _calculate_snr_diff(self, alpha_bar): + snr = alpha_bar / (1 - alpha_bar + 1e-8) + snr_prev = torch.cat([torch.tensor([0.]), snr[:-1]]) + return torch.clamp(snr - snr_prev, min=1e-5) + + def forward_denoise(self, x,y, z_prev, t): + return self.blocks[t](x,y, z_prev, self.target_embedder)[0] + + def regress(self, z): + return self.regressor(z) + + def inference(self, x, y): + B = x.size(0) + z = torch.randn(B, self.embedding_dim, device=x.device) + if not self.training: + z = torch.zeros(B, self.embedding_dim, device=x.device) + + for t in range(self.T): + z = self.forward_denoise(x,y , z, t) + + return self.regress(z) + + def forward(self, x): + if x.shape[1] >= 2: + x, y = torch.split(x, [1, x.shape[1]-1], dim=1) + else: + y = x + + if self.task=="direction": + return None, None, self.inference(x,y) + elif self.task=="energy": + return None, self.inference(x,y), None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet copy 2.py b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet copy 2.py new file mode 100644 index 00000000..f1326018 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet copy 2.py @@ -0,0 +1,116 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn1 = AdaptiveBatchNorm2d(out_channels) + self.act1 = nn.GELU() + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn2 = AdaptiveBatchNorm2d(out_channels) + self.act2 = nn.GELU() + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + def forward(self, x): + # residual = self.act1(self.bn1(self.conv1(x))) + # residual = self.act2(self.bn2(self.conv2(residual))) + residual = self.act1((self.conv1(x))) + residual = self.act2((self.conv2(residual))) + residual = self.se(residual) + out = residual + self.shortcut(x) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, drop_prob=0.2): + super().__init__() + # Ahora más profundo y ancho: + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + # nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(256, 512), # Embedding más grande + # nn.BatchNorm1d(512), + nn.GELU() + ) + + self.fc_z1 = nn.Linear(embedding_dim, 512) + self.bn_z1 = nn.BatchNorm1d(512) + self.fc_z2 = nn.Linear(512, 512) + self.bn_z2 = nn.BatchNorm1d(512) + self.fc_z3 = nn.Linear(512, 512) + self.bn_z3 = nn.BatchNorm1d(512) + self.fc_f1 = nn.Linear(1024, 512) + self.bn_f1 = nn.BatchNorm1d(512) + self.fc_f2 = nn.Linear(512, 256) + self.bn_f2 = nn.BatchNorm1d(256) + self.fc_out = nn.Linear(256, embedding_dim) + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + # h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + # h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h1 = self.act1((self.fc_z1(z_prev))) + h2 = self.act2((self.fc_z2(h1))) + + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + # h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + # h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + h_f = self.act_f1((self.fc_f1(h_f))) + h_f = self.act_f2((self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet copy.py b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet copy.py new file mode 100644 index 00000000..6487c4ee --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet copy.py @@ -0,0 +1,111 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16): + super().__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn1 = AdaptiveBatchNorm2d(out_channels) + self.act1 = nn.GELU() + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + # self.bn2 = AdaptiveBatchNorm2d(out_channels) + self.act2 = nn.GELU() + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False), + nn.BatchNorm2d(out_channels) + ) + def forward(self, x): + # residual = self.act1(self.bn1(self.conv1(x))) + # residual = self.act2(self.bn2(self.conv2(residual))) + residual = self.act1((self.conv1(x))) + residual = self.act2((self.conv2(residual))) + residual = self.se(residual) + out = residual + self.shortcut(x) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, drop_prob=0.2): + super().__init__() + # Ahora más profundo y ancho: + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 64), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(64, 128), # Más ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(128, 256), # Más profundo/ancho + nn.MaxPool2d(2), + nn.Dropout(drop_prob), + ResidualBlock(256, 256), # Otro bloque extra para profundidad + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(256, 512), # Embedding más grande + nn.BatchNorm1d(512), + nn.GELU() + ) + + self.fc_z1 = nn.Linear(embedding_dim, 512) + self.bn_z1 = nn.BatchNorm1d(512) + self.fc_z2 = nn.Linear(512, 512) + self.bn_z2 = nn.BatchNorm1d(512) + self.fc_z3 = nn.Linear(512, 512) + self.bn_z3 = nn.BatchNorm1d(512) + self.fc_f1 = nn.Linear(1024, 512) + self.bn_f1 = nn.BatchNorm1d(512) + self.fc_f2 = nn.Linear(512, 256) + self.bn_f2 = nn.BatchNorm1d(256) + self.fc_out = nn.Linear(256, embedding_dim) + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet.py b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet.py new file mode 100644 index 00000000..b74336a4 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet.py @@ -0,0 +1,170 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +class SEBlock(nn.Module): + def __init__(self, channel, reduction=16): + super(SEBlock, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.avg_pool(x).squeeze(-1).squeeze(-1) # Ensuring dimension match + y = self.fc(y).view(b, c, 1, 1) + return x * y.expand_as(x) + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_channels, out_channels, stride=1, reduction=16,use_bn=False): + + + super(BasicBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) + self.use_bn = use_bn + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) + + + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != self.expansion * out_channels: + + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(self.expansion * out_channels) + ) + + def forward(self, x): + out = F.relu((self.conv1(x))) + out = (self.conv2(out)) + + # print("out.shape", out.shape, "shortcut.shape", self.shortcut(x).shape) + # print("out.dtype", out.dtype, "shortcut.dtype", self.shortcut(x).dtype) + # print("out.device", out.device, "shortcut.device", self.shortcut(x).device) + # print("out.is_contiguous()", out.is_contiguous(), "shortcut.is_contiguous()", self.shortcut(x).is_contiguous()) + + # assert out.shape == self.shortcut(x).shape, f"Shape mismatch: {out.shape} vs {self.shortcut(x).shape}" + # assert out.dtype == self.shortcut(x).dtype, f"Dtype mismatch: {out.dtype} vs {self.shortcut(x).dtype}" + # assert out.device == self.shortcut(x).device, f"Device mismatch: {out.device} vs {self.shortcut(x).device}" + + # out = out.contiguous() + + out += self.shortcut(x).contiguous() + out = F.relu(out) + out = self.se(out) + return out + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1, block=BasicBlock, num_blocks=[2, 3, 3, 3], num_inputs=1, num_outputs=2,use_bn=False,dropout=0.0): + super().__init__() + + # block = BasicBlock + self.in_channels = 64 + self.use_bn=use_bn + + self.conv1_x = nn.Conv2d(num_channels, 64, kernel_size=3, stride=1, padding=1, bias=False) + self.conv1_y = nn.Conv2d(num_channels, 64, kernel_size=3, stride=1, padding=1, bias=False) + + self.layer1_1 = self._make_layer(block, 32, num_blocks[0], stride=1) + self.layer2_1 = self._make_layer(block, 64, num_blocks[1], stride=2) + self.layer3_1 = self._make_layer(block, 128, num_blocks[2], stride=2) + self.layer4_1 = self._make_layer(block, 256, num_blocks[3], stride=2) + + self.layer1_2 = self._make_layer(block, 32, num_blocks[0], stride=1) + self.layer2_2 = self._make_layer(block, 64, num_blocks[1], stride=2) + self.layer3_2 = self._make_layer(block, 128, num_blocks[2], stride=2) + self.layer4_2 = self._make_layer(block, 256, num_blocks[3], stride=2) + + # Reducing the number of layers and filters to make it "thin" + # self.fc_1 = nn.Linear(embedding_dim , embedding_dim) + # self.fc_2 = nn.Linear(embedding_dim , 256) + self.bn_final = nn.BatchNorm1d(512 * block.expansion) # BatchNorm layer + self.prelu = nn.PReLU(num_parameters=512 * block.expansion) # Define Leaky ReLU + + + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + self.dropout = nn.Dropout(dropout) + + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + self.fc_f1 = nn.Linear(512, 256) + self.bn_f1 = nn.BatchNorm1d(256) + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + self.fc_out = nn.Linear(128, embedding_dim) + + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + + def forward(self, x, y, z_prev, _): + + out_1 = F.relu(self.conv1_x(x)) + out_1 = self.layer1_1(out_1) + out_1 = self.layer2_1(out_1) + out_1 = self.layer3_1(out_1) + x_feat = self.layer4_1(out_1) + + out_2 = F.relu(self.conv1_y(y)) + out_2 = self.layer1_1(out_2) + out_2 = self.layer2_1(out_2) + out_2 = self.layer3_1(out_2) + y_feat = self.layer4_1(out_2) + + x_feat= self.adaptive_pool(x_feat+y_feat) + x_feat = x_feat.view(x_feat.size(0), -1) + + # h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + # h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h1 = self.act1((self.fc_z1(z_prev))) + h2 = self.act2((self.fc_z2(h1))) + + h3 = self.bn_z3(self.fc_z3(h2)) + z_feat = h3 + h1 + h_f = torch.cat([x_feat, z_feat], dim=1) + # h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + # h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + h_f = self.act_f1((self.fc_f1(h_f))) + h_f = self.act_f2((self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + return z_next, None + + def _make_layer(self, block, out_channels, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_channels, out_channels, stride, use_bn=self.use_bn)) + self.in_channels = out_channels * block.expansion + return nn.Sequential(*layers) + diff --git a/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet_original.py b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet_original.py new file mode 100644 index 00000000..fc52d51f --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/NoPropDTRegDBB/denoiseBlockThinRestNet_original.py @@ -0,0 +1,154 @@ +# Denoising block +import torch +from torch import nn +import torch.nn.functional as F + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return x * torch.sigmoid(x) + +class AdaptiveBatchNorm2d(nn.Module): + def __init__(self, num_features, eps=1e-5, momentum=0.5, affine=True): + super(AdaptiveBatchNorm2d, self).__init__() + self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine) + # self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + # self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1)) + self.a = nn.Parameter(torch.ones(1, 1, 1, 1)) + self.b = nn.Parameter(torch.zeros(1, 1, 1, 1)) + + def forward(self, x): + return self.a * x + self.b * self.bn(x) + +# class ResidualBlock(nn.Module): +# def __init__(self, in_channels, out_channels): +# super().__init__() +# self.conv_block = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels), +# nn.PReLU(), +# nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), +# # AdaptiveBatchNorm2d(out_channels) +# ) + +# self.shortcut = nn.Sequential() +# if in_channels != out_channels: +# self.shortcut = nn.Sequential( +# nn.Conv2d(in_channels, out_channels, kernel_size=1), +# nn.BatchNorm2d(out_channels) +# ) + +# self.relu = nn.PReLU() + +# def forward(self, x): +# return self.relu(self.conv_block(x) + self.shortcut(x)) + +class SEBlock(nn.Module): + def __init__(self, channels, reduction=16): + super().__init__() + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channels, channels // reduction, bias=False), + nn.ReLU(), + nn.Linear(channels // reduction, channels, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.pool(x).view(b, c) + y = self.fc(y).view(b, c, 1, 1) + return x * y + + +class ResidualBlock(nn.Module): + def __init__(self, in_channels, out_channels, reduction=16, drop_path_rate=0.1): + super().__init__() + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.norm1 = nn.GroupNorm(8, out_channels) + self.act1 = nn.GELU() + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + self.norm2 = nn.GroupNorm(8, out_channels) + self.act2 = nn.GELU() + + self.se = SEBlock(out_channels, reduction) + + self.shortcut = nn.Identity() + if in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1), + nn.GroupNorm(8, out_channels) + ) + + try: + from timm.models.layers import DropPath + self.drop_path = DropPath(drop_path_rate) + except ImportError: + self.drop_path = nn.Identity() + + # self.final_act = nn.GELU() + self.final_act = MemoryEfficientSwish() + def forward(self, x): + residual = self.act1(self.norm1(self.conv1(x))) + residual = self.act2(self.norm2(self.conv2(residual))) + # residual = self.act1((self.conv1(x))) + # residual = self.act2((self.conv2(residual))) + + residual = self.se(residual) + + out = self.drop_path(residual) + self.shortcut(x) + return self.final_act(out) + +class DenoiseBlock(nn.Module): + def __init__(self, embedding_dim, num_channels=1): + super().__init__() + + self.conv_path = nn.Sequential( + ResidualBlock(num_channels, 32), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(32, 64), + nn.MaxPool2d(2), + nn.Dropout(0.2), + ResidualBlock(64, 128), + nn.AdaptiveAvgPool2d((1, 1)), + nn.Flatten(), + nn.Linear(128, 256), + ) + + self.fc_z1 = nn.Linear(embedding_dim, 256) + self.bn_z1 = nn.BatchNorm1d(256) + self.fc_z2 = nn.Linear(256, 256) + self.bn_z2 = nn.BatchNorm1d(256) + self.fc_z3 = nn.Linear(256, 256) + self.bn_z3 = nn.BatchNorm1d(256) + + self.fc_f1 = nn.Linear(512, 256) + self.bn_f1 = nn.BatchNorm1d(256) + self.fc_f2 = nn.Linear(256, 128) + self.bn_f2 = nn.BatchNorm1d(128) + self.fc_out = nn.Linear(128, embedding_dim) + + self.act1 = nn.PReLU() + self.act2 = nn.PReLU() + self.act3 = nn.PReLU() + self.act_f1 = nn.PReLU() + self.act_f2 = nn.PReLU() + + def forward(self, x, z_prev, _): + x_feat = self.conv_path(x) + + h1 = self.act1(self.bn_z1(self.fc_z1(z_prev))) + h2 = self.act2(self.bn_z2(self.fc_z2(h1))) + h3 = self.bn_z3(self.fc_z3(h2)) + + z_feat = h3 + h1 + + h_f = torch.cat([x_feat, z_feat], dim=1) + + h_f = self.act_f1(self.bn_f1(self.fc_f1(h_f))) + h_f = self.act_f2(self.bn_f2(self.fc_f2(h_f))) + z_next = self.fc_out(h_f) + + return z_next, None diff --git a/ctlearn/core/pytorch/nets/models/PyTorchResNet/PyTorchResNet.py b/ctlearn/core/pytorch/nets/models/PyTorchResNet/PyTorchResNet.py new file mode 100644 index 00000000..1c62d349 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/PyTorchResNet/PyTorchResNet.py @@ -0,0 +1,330 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class DualSqueezeExciteBlock(nn.Module): + 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): + return self.cse(x) + self.sse(x) + +class ChannelSqueezeExciteBlock(nn.Module): + def __init__(self, in_channels, ratio=4): + super().__init__() + self.gate = nn.Sequential( + nn.Conv2d(in_channels, in_channels // ratio, kernel_size=1, bias=True), + nn.ReLU(), + nn.Conv2d(in_channels // ratio, in_channels, kernel_size=1, bias=True), + nn.Sigmoid() + ) + + def forward(self, x): + squeeze = F.adaptive_avg_pool2d(x, (1, 1)) + excitation = self.gate(squeeze) + return x * excitation + +class SpatialSqueezeExciteBlock(nn.Module): + def __init__(self, in_channels): + super().__init__() + self.spatial_conv = nn.Conv2d(in_channels, 1, kernel_size=1, bias=True) + + def forward(self, x): + spatial_mask = torch.sigmoid(self.spatial_conv(x)) + return x * spatial_mask + +class MultiHeadClassifier(nn.Module): + def __init__(self, heads_dict, task): + super().__init__() + self.heads = nn.ModuleDict(heads_dict) + self.task = task + + def forward(self, x): + if x.dim() > 2: + x = torch.flatten(x, start_dim=1) + + classification = None + energy = None + direction = None + + if self.task == "type" and "type" in self.heads: + classification = self.heads["type"](x) + if self.task == "energy" and "energy" in self.heads: + energy = self.heads["energy"](x) + if self.task == "direction" and "direction" in self.heads: + direction = self.heads["direction"](x) + + return classification, energy, direction + +def pytorch_build_fully_connect_head(in_features, layers, activation_function, tasks): + heads = {} + act_map = { + "relu": nn.ReLU, + "tanh": nn.Tanh, + "sigmoid": nn.Sigmoid + } + + for task in tasks: + if task not in layers: + continue + 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) + + task_str = tasks[0] if len(tasks) > 0 else "type" + return MultiHeadClassifier(heads, task=task_str) + +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 + + if conv_shortcut: + self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False) + else: + self.shortcut = nn.Identity() + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) + + self.setup_attention(out_channels) + + 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) + + out = F.relu(self.conv1(x)) + out = self.conv2(out) + + if self.attn_layer: + 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 + out_channels = 4 * base_filters + + if conv_shortcut: + self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False) + else: + self.shortcut = nn.Identity() + + self.conv1 = nn.Conv2d(in_channels, base_filters, kernel_size=1, stride=stride, bias=False) + self.conv2 = nn.Conv2d(base_filters, base_filters, kernel_size=3, padding=1, bias=False) + self.conv3 = nn.Conv2d(base_filters, out_channels, kernel_size=1, bias=False) + + self.setup_attention(out_channels) + + 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) + + 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 PyTorchResNetModel(nn.Module): + def __init__( + self, + task="type", + num_inputs=1, + num_outputs=2, + init_padding=0, + init_layer=None, + init_max_pool=None, + residual_block_type="bottleneck", + architecture=None, + head_layers=None, + head_activation_function=None, + attention_mechanism="Dual-SE", + attention_reduction_ratio=16, + ): + super().__init__() + self.task = task + self.init_padding = init_padding + self.init_layer = init_layer + self.init_max_pool = init_max_pool + self.residual_block_type = residual_block_type + + if architecture is None: + architecture = [ + {"filters": 48, "blocks": 2}, + {"filters": 96, "blocks": 3}, + {"filters": 128, "blocks": 3}, + {"filters": 256, "blocks": 3}, + ] + self.architecture = architecture + + if head_layers is None: + head_layers = { + "type": [512, 256, num_outputs], + "energy": [512, 256, 1], + "direction": [512, 256, 3], + } + + if head_activation_function is None: + head_activation_function = { + "type": "relu", + "energy": "relu", + "direction": "tanh", + } + + self.attention = None + if attention_mechanism is not None: + self.attention = { + "mechanism": attention_mechanism, + "reduction_ratio": attention_reduction_ratio, + } + + input_shape = (num_inputs, 224, 224) # Spatial dims don't affect init + self.backbone_model, out_features = self._build_backbone(input_shape) + + self.logits_head = pytorch_build_fully_connect_head( + out_features, head_layers, head_activation_function, [task] + ) + + def _build_backbone(self, input_shape): + in_channels = input_shape[0] + modules = [] + + if self.init_padding > 0: + modules.append(nn.ZeroPad2d(self.init_padding)) + + if self.init_layer 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=False)) + modules.append(nn.ReLU()) + in_channels = out_ch + + if self.init_max_pool 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=p_size // 2)) + + 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) + + class GlobalAvgPool(nn.Module): + def forward(self, x): + return F.adaptive_avg_pool2d(x, (1, 1)) + + modules.append(GlobalAvgPool()) + + 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] + + 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 + + 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): + block_layer = BasicBlock if residual_block_type == "basic" else BottleneckBlock + stack = [] + + base_kwargs = { + "in_channels": in_channels, + "stride": stride, + "attention": attention + } + + if residual_block_type == "basic": + base_kwargs["out_channels"] = filters + base_kwargs["base_filters"] = filters + else: + base_kwargs["base_filters"] = filters + + stack.append(block_layer(conv_shortcut=True, **base_kwargs)) + + multiplier = 4 if residual_block_type == "bottleneck" else 1 + current_in = filters * multiplier + + for _ in range(1, blocks): + next_kwargs = { + "in_channels": current_in, + "stride": 1, + "attention": attention, + "conv_shortcut": False + } + if residual_block_type == "basic": + next_kwargs["out_channels"] = filters + next_kwargs["base_filters"] = filters + else: + next_kwargs["base_filters"] = filters + + stack.append(block_layer(**next_kwargs)) + + return stack + + def forward(self, x): + features = self.backbone_model(x) + return self.logits_head(features) diff --git a/ctlearn/core/pytorch/nets/models/ResNeXtDBB/ResNeXtDBB.py b/ctlearn/core/pytorch/nets/models/ResNeXtDBB/ResNeXtDBB.py new file mode 100644 index 00000000..db693e55 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/ResNeXtDBB/ResNeXtDBB.py @@ -0,0 +1,177 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class SqueezeExcitation(nn.Module): + def __init__(self, in_channels, reduction_ratio=16): + super(SqueezeExcitation, self).__init__() + self.in_channels = in_channels + self.reduction_ratio = reduction_ratio + self.se_channels = max(in_channels // reduction_ratio, 1) # Evitar que los canales sean menos de 1 + self.squeeze = nn.AdaptiveAvgPool2d(1) + self.excitation = nn.Sequential( + nn.Linear(in_channels, self.se_channels, bias=False), + nn.ReLU(inplace=True), + nn.Linear(self.se_channels, in_channels, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + batch_size, channels, _, _ = x.size() + # Squeeze: Global Average Pooling + y = self.squeeze(x).view(batch_size, channels) + # Excitation: Dos capas densas + y = self.excitation(y).view(batch_size, channels, 1, 1) + # Recalibrar los canales + return x * y.expand_as(x) + +class ResNeXtBlock(nn.Module): + expansion = 4 # Correcto ajuste del factor de expansión + + def __init__(self, in_channels, out_channels, stride=1, groups=32, use_gn=False, reduction_ratio=16): + super(ResNeXtBlock, self).__init__() + + # Primera capa de convolución + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + self.norm1 = nn.GroupNorm(32, out_channels) if use_gn else nn.BatchNorm2d(out_channels) + + # Segunda capa de convolución con agrupación + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, padding=1, groups=groups, bias=False) + self.norm2 = nn.GroupNorm(32, out_channels) if use_gn else nn.BatchNorm2d(out_channels) + + # Tercera capa de convolución + self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, kernel_size=1, bias=False) + self.norm3 = nn.GroupNorm(32, out_channels * self.expansion) if use_gn else nn.BatchNorm2d(out_channels * self.expansion) + + # Bloque Squeeze-and-Excitation + self.se_block = SqueezeExcitation(out_channels * self.expansion, reduction_ratio) + + # Atajo (shortcut) para ajustar el número de canales si es necesario + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != out_channels * self.expansion: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels * self.expansion, kernel_size=1, stride=stride, bias=False), + nn.GroupNorm(32, out_channels * self.expansion) if use_gn else nn.BatchNorm2d(out_channels * self.expansion) + ) + self.prelu_1 = nn.PReLU() + self.prelu_2 = nn.PReLU() + self.prelu_3 = nn.PReLU() + + def forward(self, x): + # Aplicar las capas de convolución y las normalizaciones con ReLU + out = self.prelu_1(self.norm1(self.conv1(x))) + out = self.prelu_2(self.norm2(self.conv2(out))) + out = self.norm3(self.conv3(out)) + + # Apply Squeeze-and-Excitation + out = self.se_block(out) + + out += self.shortcut(x) + out = self.prelu_3(out) + return out + +class ResNeXtDBB(nn.Module): + def __init__(self,task, block, layers, num_inputs=1, num_classes=1, use_gn=False, use_concat=False, dropout_rate=0.5): + super(ResNeXtDBB, self).__init__() + self.in_channels = 64 + self.use_gn = use_gn + self.use_concat = use_concat + self.task = task + # Backbone 1 + self.conv1_a = nn.Conv2d(num_inputs, 64, kernel_size=7, stride=2, padding=3, bias=False) + self.norm1_a = nn.GroupNorm(32, 64) if use_gn else nn.BatchNorm2d(64) + + self.layer1_a = self._make_layer(block, 64, layers[0], stride=1) + self.layer2_a = self._make_layer(block, 128, layers[1], stride=2) + self.layer3_a = self._make_layer(block, 256, layers[2], stride=2) + self.layer4_a = self._make_layer(block, 512, layers[3], stride=2) + + # Backbone 2 + self.in_channels = 64 # Reset in_channels for the second backbone + self.conv1_b = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False) + self.norm1_b = nn.GroupNorm(32, 64) if use_gn else nn.BatchNorm2d(64) + + self.layer1_b = self._make_layer(block, 64, layers[0], stride=1) + self.layer2_b = self._make_layer(block, 128, layers[1], stride=2) + self.layer3_b = self._make_layer(block, 256, layers[2], stride=2) + self.layer4_b = self._make_layer(block, 512, layers[3], stride=2) + + self.prelu_1 = nn.PReLU() + self.prelu_2 = nn.PReLU() + + # Dropout layer + self.dropout = nn.Dropout(dropout_rate) + + # Fully connected layer + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + + # Ajuste de los canales para la capa completamente conectada + if self.use_concat: + self.fc = nn.Linear(512 * block.expansion * 2, num_classes) + else: + self.fc = nn.Linear(512 * block.expansion, num_classes) + + def _make_layer(self, block, out_channels, num_blocks, stride): + layers = [] + layers.append(block(self.in_channels, out_channels, stride, use_gn=self.use_gn)) + self.in_channels = out_channels * block.expansion + for _ in range(1, num_blocks): + layers.append(block(self.in_channels, out_channels, use_gn=self.use_gn)) + return nn.Sequential(*layers) + + def forward(self, x1): + if x1.shape[1] >= 2: + x1, x2 = torch.split(x1, [1, x1.shape[1]-1], dim=1) + else: + x2 = x1 + + classification=None + energy=None + direction=None + + # Backbone 1 + # out1 = F.relu(self.norm1_a(self.conv1_a(x1))) + out1 = self.prelu_1(self.conv1_a(x1)) + + out1 = self.layer1_a(out1) + out1 = self.layer2_a(out1) + out1 = self.layer3_a(out1) + out1 = self.layer4_a(out1) + out1 = self.adaptive_pool(out1) + out1 = out1.view(out1.size(0), -1) + out1 = self.dropout(out1) + + # Backbone 2 + # out2 = F.relu(self.norm1_b(self.conv1_b(x2))) + out2 = self.prelu_2(self.conv1_b(x2)) + out2 = self.layer1_b(out2) + out2 = self.layer2_b(out2) + out2 = self.layer3_b(out2) + out2 = self.layer4_b(out2) + out2 = self.adaptive_pool(out2) + out2 = out2.view(out2.size(0), -1) + out2 = self.dropout(out2) + + # Combine outputs + if self.use_concat: + out = torch.cat((out1, out2), dim=1) + else: + out = out1 + out2 + + out = self.fc(out) + + if self.task == "type": + classification = out + elif self.task == "energy": + energy = out + elif self.task == "direction": + direction = out + + return classification, energy, direction + +def ResNeXtDuo(task,num_blocks=[2, 2, 2, 2], num_inputs=1, num_classes=2, use_gn=True, use_concat=False, dropout_rate=0.5): + # Here we configure fewer blocks for a lighter model + return ResNeXtDBB(task,ResNeXtBlock, num_blocks, num_inputs, num_classes=num_classes, use_gn=use_gn, use_concat=use_concat, dropout_rate=dropout_rate) + +# Instancia del modelo +# resnext_duo = ResNeXtDuo(ResNeXtBlock, [3, 4, 6, 3], num_classes=1, use_gn=True, use_concat=True, dropout_rate=0.5) \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/StackedHGNet/StackedHGNet.py b/ctlearn/core/pytorch/nets/models/StackedHGNet/StackedHGNet.py new file mode 100644 index 00000000..e0635cb5 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/StackedHGNet/StackedHGNet.py @@ -0,0 +1,141 @@ +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# from core.coord_conv import CoordConvTh +# from lib.dataset import get_decoder + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class ConvBlock(nn.Module): + def __init__(self, inp_dim, out_dim, kernel_size=3, stride=1, bn=False, relu=True): + super().__init__() + self.conv = nn.Conv2d(inp_dim, out_dim, kernel_size, stride, padding=(kernel_size - 1) // 2, bias=not bn) + self.bn = nn.BatchNorm2d(out_dim) if bn else nn.Identity() + self.relu = nn.ReLU(inplace=True) if relu else nn.Identity() + def forward(self, x): + return self.relu(self.bn(self.conv(x))) + +class ResBlock(nn.Module): + def __init__(self, inp_dim, out_dim, mid_dim=None): + super().__init__() + mid_dim = mid_dim or out_dim // 2 + self.conv1 = ConvBlock(inp_dim, mid_dim, 1, bn=True, relu=True) + self.conv2 = ConvBlock(mid_dim, mid_dim, 3, bn=True, relu=True) + self.conv3 = ConvBlock(mid_dim, out_dim, 1, bn=True, relu=False) + self.skip = ConvBlock(inp_dim, out_dim, 1, bn=True, relu=False) if inp_dim != out_dim else nn.Identity() + self.relu = nn.ReLU(inplace=True) + def forward(self, x): + out = self.conv1(x) + out = self.conv2(out) + out = self.conv3(out) + return self.relu(out + self.skip(x)) + +class Hourglass(nn.Module): + def __init__(self, n, f): + super().__init__() + self.up1 = ResBlock(f, f) + self.pool1 = nn.MaxPool2d(2, 2) + self.low1 = ResBlock(f, f) + if n > 1: + self.low2 = Hourglass(n - 1, f) + else: + self.low2 = ResBlock(f, f) + self.low3 = ResBlock(f, f) + self.up2 = nn.Upsample(scale_factor=2, mode='nearest') + def forward(self, x): + up1 = self.up1(x) + low1 = self.low1(self.pool1(x)) + low2 = self.low2(low1) + low3 = self.low3(low2) + up2 = self.up2(low3) + if up2.shape[-2:] != up1.shape[-2:]: + up2 = F.interpolate(up2, size=up1.shape[-2:], mode='nearest') + return up1 + up2 + +class StackedHGNet(nn.Module): + def __init__(self,task, input_channels=3, nstack=2, nlevels=4, in_channel=256, output_dim=3, use_bn=True, use_stn=False): + super().__init__() + self.task = task + self.nstack = nstack + self.pre = nn.Sequential( + ConvBlock(input_channels, 64, 7, 2, bn=use_bn, relu=True), + ResBlock(64, 128), + nn.MaxPool2d(2, 2), + ResBlock(128, 128), + ResBlock(128, in_channel) + ) + self.hgs = nn.ModuleList([ + Hourglass(nlevels, in_channel) for _ in range(nstack) + ]) + self.features = nn.ModuleList([ + nn.Sequential( + ResBlock(in_channel, in_channel), + ConvBlock(in_channel, in_channel, 1, bn=use_bn, relu=True) + ) for _ in range(nstack) + ]) + # self.out_regression = nn.ModuleList([ + # nn.Sequential( + # nn.AdaptiveAvgPool2d(1), + # nn.Flatten(), + # nn.Linear(in_channel, output_dim) + # ) for _ in range(nstack) + # ]) + + self.out_regression = nn.ModuleList([ + nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + nn.Linear(in_channel, 128), + nn.PReLU(), + nn.Linear(128, 64), + nn.PReLU(), + nn.Linear(64, output_dim) + ) for _ in range(nstack) + ]) + + # self.out_regression = nn.ModuleList([ + # nn.Sequential( + # nn.AdaptiveAvgPool2d(1), + # nn.Flatten(), + # nn.Linear(in_channel, 128), + # nn.ReLU(), + # nn.BatchNorm1d(128), + # nn.Linear(128, 64), + # nn.ReLU(), + # nn.Linear(64, output_dim) + # ) for _ in range(nstack) + # ]) + # Head auxiliar de heatmap (un canal) + # self.out_heatmap = nn.ModuleList([ + # ConvBlock(in_channel, 1, 1, relu=False, bn=False) for _ in range(nstack) + # ]) + self.merge_features = nn.ModuleList([ + ConvBlock(in_channel, in_channel, 1, relu=False, bn=False) + for _ in range(nstack - 1) + ]) + + def forward(self, x): + x = self.pre(x) + regression_outputs = [] + # heatmap_outputs = [] + for i in range(self.nstack): + hg = self.hgs[i](x) + feature = self.features[i](hg) + regression_output = self.out_regression[i](feature) + # heatmap_output = self.out_heatmap[i](feature) + regression_outputs.append(regression_output) + # heatmap_outputs.append(heatmap_output) + if i < self.nstack - 1: + x = x + self.merge_features[i](feature) + + if self.task=="direction": + return None, None, regression_outputs[-1] + elif self.task=="energy": + return None, regression_outputs[-1], None + else: + raise ValueError(f"No implemented") diff --git a/ctlearn/core/pytorch/nets/models/StackedHGNet/core/coord_conv.py b/ctlearn/core/pytorch/nets/models/StackedHGNet/core/coord_conv.py new file mode 100644 index 00000000..7239421d --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/StackedHGNet/core/coord_conv.py @@ -0,0 +1,157 @@ +import torch +import torch.nn as nn + + +class AddCoordsTh(nn.Module): + def __init__(self, x_dim, y_dim, with_r=False, with_boundary=False): + super(AddCoordsTh, self).__init__() + self.x_dim = x_dim + self.y_dim = y_dim + self.with_r = with_r + self.with_boundary = with_boundary + + def forward(self, input_tensor, heatmap=None): + """ + input_tensor: (batch, c, x_dim, y_dim) + """ + batch_size_tensor = input_tensor.shape[0] + + xx_ones = torch.ones([1, self.y_dim], dtype=torch.int32).to(input_tensor) + xx_ones = xx_ones.unsqueeze(-1) + + xx_range = torch.arange(self.x_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor) + xx_range = xx_range.unsqueeze(1) + + xx_channel = torch.matmul(xx_ones.float(), xx_range.float()) + xx_channel = xx_channel.unsqueeze(-1) + + yy_ones = torch.ones([1, self.x_dim], dtype=torch.int32).to(input_tensor) + yy_ones = yy_ones.unsqueeze(1) + + yy_range = torch.arange(self.y_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor) + yy_range = yy_range.unsqueeze(-1) + + yy_channel = torch.matmul(yy_range.float(), yy_ones.float()) + yy_channel = yy_channel.unsqueeze(-1) + + xx_channel = xx_channel.permute(0, 3, 2, 1) + yy_channel = yy_channel.permute(0, 3, 2, 1) + + xx_channel = xx_channel / (self.x_dim - 1) + yy_channel = yy_channel / (self.y_dim - 1) + + xx_channel = xx_channel * 2 - 1 + yy_channel = yy_channel * 2 - 1 + + xx_channel = xx_channel.repeat(batch_size_tensor, 1, 1, 1) + yy_channel = yy_channel.repeat(batch_size_tensor, 1, 1, 1) + + if self.with_boundary and type(heatmap) != type(None): + boundary_channel = torch.clamp(heatmap[:, -1:, :, :], + 0.0, 1.0) + + zero_tensor = torch.zeros_like(xx_channel).to(xx_channel) + xx_boundary_channel = torch.where(boundary_channel>0.05, + xx_channel, zero_tensor) + yy_boundary_channel = torch.where(boundary_channel>0.05, + yy_channel, zero_tensor) + ret = torch.cat([input_tensor, xx_channel, yy_channel], dim=1) + + + if self.with_r: + rr = torch.sqrt(torch.pow(xx_channel, 2) + torch.pow(yy_channel, 2)) + rr = rr / torch.max(rr) + ret = torch.cat([ret, rr], dim=1) + + if self.with_boundary and type(heatmap) != type(None): + ret = torch.cat([ret, xx_boundary_channel, + yy_boundary_channel], dim=1) + return ret + + +class CoordConvTh(nn.Module): + """CoordConv layer as in the paper.""" + def __init__(self, x_dim, y_dim, with_r, with_boundary, + in_channels, out_channels, first_one=False, relu=False, bn=False, *args, **kwargs): + super(CoordConvTh, self).__init__() + self.addcoords = AddCoordsTh(x_dim=x_dim, y_dim=y_dim, with_r=with_r, + with_boundary=with_boundary) + in_channels += 2 + if with_r: + in_channels += 1 + if with_boundary and not first_one: + in_channels += 2 + self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, *args, **kwargs) + self.relu = nn.ReLU() if relu else None + self.bn = nn.BatchNorm2d(out_channels) if bn else None + + self.with_boundary = with_boundary + self.first_one = first_one + + + def forward(self, input_tensor, heatmap=None): + assert (self.with_boundary and not self.first_one) == (heatmap is not None) + ret = self.addcoords(input_tensor, heatmap) + ret = self.conv(ret) + if self.bn is not None: + ret = self.bn(ret) + if self.relu is not None: + ret = self.relu(ret) + + return ret + + +''' +An alternative implementation for PyTorch with auto-infering the x-y dimensions. +''' +class AddCoords(nn.Module): + + def __init__(self, with_r=False): + super().__init__() + self.with_r = with_r + + def forward(self, input_tensor): + """ + Args: + input_tensor: shape(batch, channel, x_dim, y_dim) + """ + batch_size, _, x_dim, y_dim = input_tensor.size() + + xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1).to(input_tensor) + yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2).to(input_tensor) + + xx_channel = xx_channel / (x_dim - 1) + yy_channel = yy_channel / (y_dim - 1) + + xx_channel = xx_channel * 2 - 1 + yy_channel = yy_channel * 2 - 1 + + xx_channel = xx_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3) + yy_channel = yy_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3) + + ret = torch.cat([ + input_tensor, + xx_channel.type_as(input_tensor), + yy_channel.type_as(input_tensor)], dim=1) + + if self.with_r: + rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2)) + ret = torch.cat([ret, rr], dim=1) + + return ret + + +class CoordConv(nn.Module): + + def __init__(self, in_channels, out_channels, with_r=False, **kwargs): + super().__init__() + self.addcoords = AddCoords(with_r=with_r) + in_channels += 2 + if with_r: + in_channels += 1 + self.conv = nn.Conv2d(in_channels, out_channels, **kwargs) + + def forward(self, x): + ret = self.addcoords(x) + ret = self.conv(ret) + return ret diff --git a/ctlearn/core/pytorch/nets/models/StackedHGNetDBB/StackedHGNetDBB.py b/ctlearn/core/pytorch/nets/models/StackedHGNetDBB/StackedHGNetDBB.py new file mode 100644 index 00000000..69a96fea --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/StackedHGNetDBB/StackedHGNetDBB.py @@ -0,0 +1,151 @@ +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# from core.coord_conv import CoordConvTh +# from lib.dataset import get_decoder + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class ConvBlock(nn.Module): + def __init__(self, inp_dim, out_dim, kernel_size=3, stride=1, bn=False, relu=True): + super().__init__() + self.conv = nn.Conv2d(inp_dim, out_dim, kernel_size, stride, padding=(kernel_size - 1) // 2, bias=not bn) + self.bn = nn.BatchNorm2d(out_dim) if bn else nn.Identity() + self.relu = nn.ReLU(inplace=True) if relu else nn.Identity() + def forward(self, x): + return self.relu(self.bn(self.conv(x))) + +class ResBlock(nn.Module): + def __init__(self, inp_dim, out_dim, mid_dim=None): + super().__init__() + mid_dim = mid_dim or out_dim // 2 + self.conv1 = ConvBlock(inp_dim, mid_dim, 1, bn=True, relu=True) + self.conv2 = ConvBlock(mid_dim, mid_dim, 3, bn=True, relu=True) + self.conv3 = ConvBlock(mid_dim, out_dim, 1, bn=True, relu=False) + self.skip = ConvBlock(inp_dim, out_dim, 1, bn=True, relu=False) if inp_dim != out_dim else nn.Identity() + self.relu = nn.ReLU(inplace=True) + def forward(self, x): + out = self.conv1(x) + out = self.conv2(out) + out = self.conv3(out) + return self.relu(out + self.skip(x)) + +class Hourglass(nn.Module): + def __init__(self, n, f): + super().__init__() + self.up1 = ResBlock(f, f) + self.pool1 = nn.MaxPool2d(2, 2) + self.low1 = ResBlock(f, f) + if n > 1: + self.low2 = Hourglass(n - 1, f) + else: + self.low2 = ResBlock(f, f) + self.low3 = ResBlock(f, f) + self.up2 = nn.Upsample(scale_factor=2, mode='nearest') + def forward(self, x): + up1 = self.up1(x) + low1 = self.low1(self.pool1(x)) + low2 = self.low2(low1) + low3 = self.low3(low2) + up2 = self.up2(low3) + if up2.shape[-2:] != up1.shape[-2:]: + up2 = F.interpolate(up2, size=up1.shape[-2:], mode='nearest') + return up1 + up2 + +class StackedHGNetDBB(nn.Module): + def __init__(self,task, input_channels=3, nstack=2, nlevels=4, in_channel=256, output_dim=3, use_bn=True, use_stn=False): + super().__init__() + self.task = task + self.nstack = nstack + self.pre_x = nn.Sequential( + ConvBlock(input_channels, 64, 7, 2, bn=use_bn, relu=True), + ResBlock(64, 128), + nn.MaxPool2d(2, 2), + ResBlock(128, 128), + ResBlock(128, in_channel) + ) + self.pre_y = nn.Sequential( + ConvBlock(input_channels, 64, 7, 2, bn=use_bn, relu=True), + ResBlock(64, 128), + nn.MaxPool2d(2, 2), + ResBlock(128, 128), + ResBlock(128, in_channel) + ) + self.hgs = nn.ModuleList([ + Hourglass(nlevels, in_channel) for _ in range(nstack) + ]) + self.features = nn.ModuleList([ + nn.Sequential( + ResBlock(in_channel, in_channel), + ConvBlock(in_channel, in_channel, 1, bn=use_bn, relu=True) + ) for _ in range(nstack) + ]) + # self.out_regression = nn.ModuleList([ + # nn.Sequential( + # nn.AdaptiveAvgPool2d(1), + # nn.Flatten(), + # nn.Linear(in_channel, output_dim) + # ) for _ in range(nstack) + # ]) + + self.out_regression = nn.ModuleList([ + nn.Sequential( + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + nn.Linear(in_channel, 128), + nn.PReLU(), + nn.Linear(128, 64), + nn.PReLU(), + nn.Linear(64, output_dim) + ) for _ in range(nstack) + ]) + + # self.out_regression = nn.ModuleList([ + # nn.Sequential( + # nn.AdaptiveAvgPool2d(1), + # nn.Flatten(), + # nn.Linear(in_channel, 128), + # nn.ReLU(), + # nn.BatchNorm1d(128), + # nn.Linear(128, 64), + # nn.ReLU(), + # nn.Linear(64, output_dim) + # ) for _ in range(nstack) + # ]) + # Head auxiliar de heatmap (un canal) + # self.out_heatmap = nn.ModuleList([ + # ConvBlock(in_channel, 1, 1, relu=False, bn=False) for _ in range(nstack) + # ]) + self.merge_features = nn.ModuleList([ + ConvBlock(in_channel, in_channel, 1, relu=False, bn=False) + for _ in range(nstack - 1) + ]) + + def forward(self, x, y ): + x = self.pre_x(x) + y = self.pre_y(y) + + x = x + y + regression_outputs = [] + # heatmap_outputs = [] + for i in range(self.nstack): + hg = self.hgs[i](x) + feature = self.features[i](hg) + regression_output = self.out_regression[i](feature) + # heatmap_output = self.out_heatmap[i](feature) + regression_outputs.append(regression_output) + # heatmap_outputs.append(heatmap_output) + if i < self.nstack - 1: + x = x + self.merge_features[i](feature) + + if self.task=="direction": + return None, None, regression_outputs[-1] + elif self.task=="energy": + return None, regression_outputs[-1], None + else: + raise ValueError(f"No implemented") \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/StackedHGNetDBB/core/coord_conv.py b/ctlearn/core/pytorch/nets/models/StackedHGNetDBB/core/coord_conv.py new file mode 100644 index 00000000..7239421d --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/StackedHGNetDBB/core/coord_conv.py @@ -0,0 +1,157 @@ +import torch +import torch.nn as nn + + +class AddCoordsTh(nn.Module): + def __init__(self, x_dim, y_dim, with_r=False, with_boundary=False): + super(AddCoordsTh, self).__init__() + self.x_dim = x_dim + self.y_dim = y_dim + self.with_r = with_r + self.with_boundary = with_boundary + + def forward(self, input_tensor, heatmap=None): + """ + input_tensor: (batch, c, x_dim, y_dim) + """ + batch_size_tensor = input_tensor.shape[0] + + xx_ones = torch.ones([1, self.y_dim], dtype=torch.int32).to(input_tensor) + xx_ones = xx_ones.unsqueeze(-1) + + xx_range = torch.arange(self.x_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor) + xx_range = xx_range.unsqueeze(1) + + xx_channel = torch.matmul(xx_ones.float(), xx_range.float()) + xx_channel = xx_channel.unsqueeze(-1) + + yy_ones = torch.ones([1, self.x_dim], dtype=torch.int32).to(input_tensor) + yy_ones = yy_ones.unsqueeze(1) + + yy_range = torch.arange(self.y_dim, dtype=torch.int32).unsqueeze(0).to(input_tensor) + yy_range = yy_range.unsqueeze(-1) + + yy_channel = torch.matmul(yy_range.float(), yy_ones.float()) + yy_channel = yy_channel.unsqueeze(-1) + + xx_channel = xx_channel.permute(0, 3, 2, 1) + yy_channel = yy_channel.permute(0, 3, 2, 1) + + xx_channel = xx_channel / (self.x_dim - 1) + yy_channel = yy_channel / (self.y_dim - 1) + + xx_channel = xx_channel * 2 - 1 + yy_channel = yy_channel * 2 - 1 + + xx_channel = xx_channel.repeat(batch_size_tensor, 1, 1, 1) + yy_channel = yy_channel.repeat(batch_size_tensor, 1, 1, 1) + + if self.with_boundary and type(heatmap) != type(None): + boundary_channel = torch.clamp(heatmap[:, -1:, :, :], + 0.0, 1.0) + + zero_tensor = torch.zeros_like(xx_channel).to(xx_channel) + xx_boundary_channel = torch.where(boundary_channel>0.05, + xx_channel, zero_tensor) + yy_boundary_channel = torch.where(boundary_channel>0.05, + yy_channel, zero_tensor) + ret = torch.cat([input_tensor, xx_channel, yy_channel], dim=1) + + + if self.with_r: + rr = torch.sqrt(torch.pow(xx_channel, 2) + torch.pow(yy_channel, 2)) + rr = rr / torch.max(rr) + ret = torch.cat([ret, rr], dim=1) + + if self.with_boundary and type(heatmap) != type(None): + ret = torch.cat([ret, xx_boundary_channel, + yy_boundary_channel], dim=1) + return ret + + +class CoordConvTh(nn.Module): + """CoordConv layer as in the paper.""" + def __init__(self, x_dim, y_dim, with_r, with_boundary, + in_channels, out_channels, first_one=False, relu=False, bn=False, *args, **kwargs): + super(CoordConvTh, self).__init__() + self.addcoords = AddCoordsTh(x_dim=x_dim, y_dim=y_dim, with_r=with_r, + with_boundary=with_boundary) + in_channels += 2 + if with_r: + in_channels += 1 + if with_boundary and not first_one: + in_channels += 2 + self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, *args, **kwargs) + self.relu = nn.ReLU() if relu else None + self.bn = nn.BatchNorm2d(out_channels) if bn else None + + self.with_boundary = with_boundary + self.first_one = first_one + + + def forward(self, input_tensor, heatmap=None): + assert (self.with_boundary and not self.first_one) == (heatmap is not None) + ret = self.addcoords(input_tensor, heatmap) + ret = self.conv(ret) + if self.bn is not None: + ret = self.bn(ret) + if self.relu is not None: + ret = self.relu(ret) + + return ret + + +''' +An alternative implementation for PyTorch with auto-infering the x-y dimensions. +''' +class AddCoords(nn.Module): + + def __init__(self, with_r=False): + super().__init__() + self.with_r = with_r + + def forward(self, input_tensor): + """ + Args: + input_tensor: shape(batch, channel, x_dim, y_dim) + """ + batch_size, _, x_dim, y_dim = input_tensor.size() + + xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1).to(input_tensor) + yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1).transpose(1, 2).to(input_tensor) + + xx_channel = xx_channel / (x_dim - 1) + yy_channel = yy_channel / (y_dim - 1) + + xx_channel = xx_channel * 2 - 1 + yy_channel = yy_channel * 2 - 1 + + xx_channel = xx_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3) + yy_channel = yy_channel.repeat(batch_size, 1, 1, 1).transpose(2, 3) + + ret = torch.cat([ + input_tensor, + xx_channel.type_as(input_tensor), + yy_channel.type_as(input_tensor)], dim=1) + + if self.with_r: + rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2)) + ret = torch.cat([ret, rr], dim=1) + + return ret + + +class CoordConv(nn.Module): + + def __init__(self, in_channels, out_channels, with_r=False, **kwargs): + super().__init__() + self.addcoords = AddCoords(with_r=with_r) + in_channels += 2 + if with_r: + in_channels += 1 + self.conv = nn.Conv2d(in_channels, out_channels, **kwargs) + + def forward(self, x): + ret = self.addcoords(x) + ret = self.conv(ret) + return ret diff --git a/ctlearn/core/pytorch/nets/models/ThinResNet/ThinResNet.py b/ctlearn/core/pytorch/nets/models/ThinResNet/ThinResNet.py new file mode 100644 index 00000000..b1b32e23 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/ThinResNet/ThinResNet.py @@ -0,0 +1,181 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from ctlearn.core.pytorch.nets.block.cnn_blocks import NormalInvGamma + +class SEBlock(nn.Module): + def __init__(self, channel, reduction=16): + super(SEBlock, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.avg_pool(x).squeeze(-1).squeeze(-1) # Ensuring dimension match + y = self.fc(y).view(b, c, 1, 1) + return x * y.expand_as(x) + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_channels, out_channels, stride=1, reduction=16,use_bn=True): + + + super(BasicBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) + self.use_bn = use_bn + if self.use_bn: + self.bn1 = nn.BatchNorm2d(out_channels) + else: + self.bn1 = nn.Identity() + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) + if self.use_bn: + self.bn2 = nn.BatchNorm2d(out_channels) + else: + self.bn2 = nn.Identity() + + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != self.expansion * out_channels: + + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(self.expansion * out_channels) + ) + + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out += self.shortcut(x) + out = F.relu(out) + out = self.se(out) + return out + +class ThinResNet(nn.Module): + def __init__(self,task, block=BasicBlock, num_blocks=[2, 3, 3, 3], num_inputs=1, num_outputs=2,use_bn=False,dropout=0.0): + super(ThinResNet, self).__init__() + + # block = BasicBlock + self.in_channels = 64 + self.use_bn=use_bn + self.task = task + self.conv1 = nn.Conv2d(num_inputs, 64, kernel_size=3, stride=1, padding=1, bias=False) + if self.use_bn: + self.bn1 = nn.BatchNorm2d(64) + else: + self.bn1 = nn.Identity() + + self.layer1_1 = self._make_layer(block, 64, num_blocks[0], stride=1) + self.layer2_1 = self._make_layer(block, 128, num_blocks[1], stride=2) + self.layer3_1 = self._make_layer(block, 256, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2) + # Reducing the number of layers and filters to make it "thin" + self.fc_1 = nn.Linear(512 * block.expansion, 512 * block.expansion) + self.fc_2 = nn.Linear(512 * block.expansion, num_outputs) + self.bn_final = nn.BatchNorm1d(512 * block.expansion) # BatchNorm layer + self.prelu = nn.PReLU(num_parameters=512 * block.expansion) # Define Leaky ReLU + if self.task == "direction": + self.normal_inv = NormalInvGamma(512 * block.expansion,num_outputs) + + if self.use_bn: + self.bn2 = nn.BatchNorm2d(64) + else: + self.bn2 = nn.Identity() + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + self.dropout = nn.Dropout(dropout) + def _make_layer(self, block, out_channels, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_channels, out_channels, stride, use_bn=self.use_bn)) + self.in_channels = out_channels * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + energy = None + classification = None + direction = None + + out = F.relu(self.conv1(x)) + out = self.layer1_1(out) + out = self.layer2_1(out) + out = self.layer3_1(out) + + out = self.layer4(out) + out = self.adaptive_pool(out) + out_feature = out.view(out.size(0), -1) + out = self.dropout(out_feature) + + # if self.training: + # out_sep = self.fc_1_separation(out) + # out_sep = self.fc_2_separation(out_sep) + + out = self.fc_1(out) + # out = self.bn_final(out) + # out = self.prelu(out) + # Original + # out = self.fc_2(out) + + + if self.task == "type": + out = self.fc_2(out) + classification = out + if self.task == "energy": + out = self.fc_2(out) + # energy = [out, out_feature] + energy = out + + if self.task == "direction": + direction = self.normal_inv(out) + + # direction = [direction, out_feature] + # if self.training: + # out = self.normal_inv(out) + # direction = out + # else: + # direction = self.normal_inv(out) + + + # if self.training: + # out = torch.cat((out, out_sep), dim=1) + + return classification, energy, direction + +# def thin_resnet34(num_blocks=[2, 3, 3, 3], num_inputs=1, num_classes=2): +# # Here we configure fewer blocks for a lighter model +# return ThinResNet_DBB(BasicBlock, num_blocks,num_inputs,num_classes) + +# def create_model(num_blocks=[2, 3, 3, 3], num_inputs=1, num_classes=2, use_bn=False, dropout=0.0): +# return ThinResNet_DBB( +# block=BasicBlock, +# num_blocks=num_blocks, +# num_inputs=num_inputs, +# num_classes=num_classes, +# use_bn=use_bn, +# dropout=dropout +# ) + +# model = thin_resnet34() +# print(model) + +# # Set the model to evaluation mode (as we are just testing with a forward pass) +# model.eval() + +# # Create dummy input tensors +# # Assuming the input images are 224x224 pixels with 1 input channel (grayscale) +# x = torch.randn(1, 1, 224, 224) # Batch size of 1 +# y = torch.randn(1, 1, 224, 224) # Batch size of 1 + +# # Forward pass through the model +# with torch.no_grad(): # We don't need to calculate gradients here +# output = model(x, y) + +# # Print the output tensor +# print("Output:", output) \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/ThinResNet_DBB/ThinResNet_DBB.py b/ctlearn/core/pytorch/nets/models/ThinResNet_DBB/ThinResNet_DBB.py new file mode 100644 index 00000000..e096e4bf --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/ThinResNet_DBB/ThinResNet_DBB.py @@ -0,0 +1,215 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from ctlearn.core.pytorch.nets.block.cnn_blocks import NormalInvGamma + +class SEBlock(nn.Module): + def __init__(self, channel, reduction=16): + super(SEBlock, self).__init__() + self.avg_pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Linear(channel, channel // reduction, bias=False), + nn.ReLU(inplace=True), + nn.Linear(channel // reduction, channel, bias=False), + nn.Sigmoid() + ) + + def forward(self, x): + b, c, _, _ = x.size() + y = self.avg_pool(x).squeeze(-1).squeeze(-1) # Ensuring dimension match + y = self.fc(y).view(b, c, 1, 1) + return x * y.expand_as(x) + +class BasicBlock(nn.Module): + expansion = 1 + + def __init__(self, in_channels, out_channels, stride=1, reduction=16,use_bn=True): + + + super(BasicBlock, self).__init__() + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False) + self.use_bn = use_bn + if self.use_bn: + self.bn1 = nn.BatchNorm2d(out_channels) + else: + self.bn1 = nn.Identity() + + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False) + if self.use_bn: + self.bn2 = nn.BatchNorm2d(out_channels) + else: + self.bn2 = nn.Identity() + + self.se = SEBlock(out_channels, reduction) + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != self.expansion * out_channels: + + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(self.expansion * out_channels) + ) + + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out += self.shortcut(x) + out = F.relu(out) + out = self.se(out) + return out + +class ThinResNet_DBB(nn.Module): + def __init__(self,task, block=BasicBlock, num_blocks=[2, 3, 3, 3], num_inputs=1, num_outputs=2,use_bn=False,dropout=0.0,extract_uncertenty=False): + super(ThinResNet_DBB, self).__init__() + self.extract_uncertenty = extract_uncertenty + # block = BasicBlock + self.in_channels = 64 + self.use_bn=use_bn + self.task = task + self.conv1 = nn.Conv2d(num_inputs, 64, kernel_size=3, stride=1, padding=1, bias=False) + if self.use_bn: + self.bn1 = nn.BatchNorm2d(64) + else: + self.bn1 = nn.Identity() + + self.layer1_1 = self._make_layer(block, 64, num_blocks[0], stride=1) + self.layer2_1 = self._make_layer(block, 128, num_blocks[1], stride=2) + self.layer3_1 = self._make_layer(block, 256, num_blocks[2], stride=2) + self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2) + # Reducing the number of layers and filters to make it "thin" + self.fc_1 = nn.Linear(512 * block.expansion, 512 * block.expansion) + self.fc_2 = nn.Linear(512 * block.expansion, num_outputs) + self.bn_final = nn.BatchNorm1d(512 * block.expansion) # BatchNorm layer + self.prelu = nn.PReLU(num_parameters=512 * block.expansion) # Define Leaky ReLU + if self.task == "direction": + self.normal_inv = NormalInvGamma(512 * block.expansion,num_outputs) + + # self.fc_1_separation = nn.Linear(512 * block.expansion, 512 * block.expansion) + # self.fc_2_separation = nn.Linear(512 * block.expansion, 4) + + self.in_channels = 64 + self.conv2 = nn.Conv2d(num_inputs, 64, kernel_size=3, stride=1, padding=1, bias=False) + self.layer1_2 = self._make_layer(block, 64, num_blocks[0], stride=1) + self.layer2_2 = self._make_layer(block, 128, num_blocks[1], stride=2) + self.layer3_2 = self._make_layer(block, 256, num_blocks[3], stride=2) # Change for 2 or 3 + #self.layer3_2 = self._make_layer(block, 256, num_blocks[3], stride=2) + + + if self.use_bn: + self.bn2 = nn.BatchNorm2d(64) + else: + self.bn2 = nn.Identity() + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + self.dropout = nn.Dropout(dropout) + def _make_layer(self, block, out_channels, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_channels, out_channels, stride, use_bn=self.use_bn)) + self.in_channels = out_channels * block.expansion + return nn.Sequential(*layers) + + def extract_feature_vector(self,x, y): + # out_1 = F.relu(self.bn1(self.conv1(x))) + out_1 = F.relu(self.conv1(x)) + out_1 = self.layer1_1(out_1) + out_1 = self.layer2_1(out_1) + out_1 = self.layer3_1(out_1) + + # out_2 = F.relu(self.bn2(self.conv2(y))) + out_2 = F.relu((self.conv2(y))) + out_2 = self.layer1_2(out_2) + out_2 = self.layer2_2(out_2) + out_2 = self.layer3_2(out_2) + out = out_1 + out_2 + + # out = self.layer3(out) + out = self.layer4(out) + out = self.adaptive_pool(out) + out_feature = out.view(out.size(0), -1) + fused_features = self.dropout(out_feature) + return fused_features + + def forward(self, x): + if x.shape[1] >= 2: + x, y = torch.split(x, [1, x.shape[1]-1], dim=1) + else: + y = x + + energy = None + classification = None + direction = None + + fused_features = self.extract_feature_vector(x,y) + # out = self.layer3(out) + + # if self.training: + # out_sep = self.fc_1_separation(out) + # out_sep = self.fc_2_separation(out_sep) + + out = self.fc_1(fused_features) + # out = self.bn_final(out) + # out = self.prelu(out) + # Original + # out = self.fc_2(out) + + + if self.task == "type": + out = self.fc_2(out) + classification = out + if self.task == "energy": + out = self.fc_2(out) + # energy = [out, out_feature] + energy = out + + if self.task == "direction": + if self.extract_uncertenty: + direction = self.normal_inv(out) + else: + direction = self.fc_2(out) + + + # direction = [direction, out_feature] + # if self.training: + # out = self.normal_inv(out) + # direction = out + # else: + # direction = self.normal_inv(out) + + + # if self.training: + # out = torch.cat((out, out_sep), dim=1) + + return [classification,fused_features], [energy,fused_features], [direction,fused_features] + +# def thin_resnet34(num_blocks=[2, 3, 3, 3], num_inputs=1, num_classes=2): +# # Here we configure fewer blocks for a lighter model +# return ThinResNet_DBB(BasicBlock, num_blocks,num_inputs,num_classes) + +# def create_model(num_blocks=[2, 3, 3, 3], num_inputs=1, num_classes=2, use_bn=False, dropout=0.0): +# return ThinResNet_DBB( +# block=BasicBlock, +# num_blocks=num_blocks, +# num_inputs=num_inputs, +# num_classes=num_classes, +# use_bn=use_bn, +# dropout=dropout +# ) + +# model = thin_resnet34() +# print(model) + +# # Set the model to evaluation mode (as we are just testing with a forward pass) +# model.eval() + +# # Create dummy input tensors +# # Assuming the input images are 224x224 pixels with 1 input channel (grayscale) +# x = torch.randn(1, 1, 224, 224) # Batch size of 1 +# y = torch.randn(1, 1, 224, 224) # Batch size of 1 + +# # Forward pass through the model +# with torch.no_grad(): # We don't need to calculate gradients here +# output = model(x, y) + +# # Print the output tensor +# print("Output:", output) \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/Transformers/TransformerDuo.py b/ctlearn/core/pytorch/nets/models/Transformers/TransformerDuo.py new file mode 100644 index 00000000..5c506ca2 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/Transformers/TransformerDuo.py @@ -0,0 +1,409 @@ +""" +Dual-Backbone Transformer Network Module + +This module implements a dual-backbone architecture combining Convolutional Neural Networks +with Transformer encoder blocks for processing Cherenkov telescope data. The architecture +uses bottleneck transformer blocks that integrate self-attention mechanisms into a +convolutional backbone, providing both local feature extraction and global context modeling. + +The dual-backbone design allows independent processing of charge and timing information +from telescope cameras before fusion for final predictions. + +Classes: + BottleneckTransformerBlock: Bottleneck block with transformer encoder layer + TransformerDBB: Dual-backbone transformer network + TransformerDuo: Factory function for creating TransformerDBB models + +References: + - "Attention Is All You Need" (Vaswani et al., NeurIPS 2017) + - "BoTNet: Bottleneck Transformers for Visual Recognition" (Srinivas et al., CVPR 2021) +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class BottleneckTransformerBlock(nn.Module): + """ + Bottleneck block with integrated transformer encoder layer. + + This block combines the efficiency of bottleneck architectures with the + global receptive field of transformers. It uses dimensionality reduction + before applying self-attention, making it computationally efficient while + still capturing long-range dependencies. + + Architecture: + Input → Conv1x1(reduce) → Norm → ReLU → + Transformer(self-attention) → + Conv1x1(expand) → Norm → (+) Shortcut → ReLU → Output + + The transformer block processes spatial features as a sequence, allowing + each position to attend to all other positions, capturing global context + that convolutional layers might miss. + + Attributes: + expansion (int): Channel expansion factor (set to 1) + conv1 (nn.Conv2d): 1x1 conv for channel reduction + norm1 (nn.BatchNorm2d or nn.GroupNorm): Normalization after reduction + transformer_block (nn.TransformerEncoderLayer): Self-attention layer + conv2 (nn.Conv2d): 1x1 conv for channel expansion + norm2 (nn.BatchNorm2d or nn.GroupNorm): Normalization after expansion + shortcut (nn.Sequential): Skip connection (identity or 1x1 conv) + """ + + expansion = 1 + + def __init__(self, in_channels, out_channels, stride=1, reduction=4, use_gn=False): + """ + Initialize the bottleneck transformer block. + + Args: + in_channels (int): Number of input channels + out_channels (int): Number of output channels + stride (int, optional): Stride for the shortcut connection. Defaults to 1 + If stride > 1, downsampling is applied in the shortcut + reduction (int, optional): Channel reduction factor. Defaults to 4 + The intermediate dimension is out_channels // reduction + Higher reduction = fewer parameters but may lose information + use_gn (bool, optional): Whether to use GroupNorm instead of BatchNorm. + Defaults to False + GroupNorm is more stable for small batch sizes + """ + super(BottleneckTransformerBlock, self).__init__() + + # Channel reduction: compress features before transformer + # Reduces computational cost of self-attention + self.conv1 = nn.Conv2d( + in_channels, + out_channels // reduction, + kernel_size=1, + stride=1, + bias=False + ) + if use_gn: + self.norm1 = nn.GroupNorm(32, out_channels // reduction) + else: + self.norm1 = nn.BatchNorm2d(out_channels // reduction) + + # Transformer encoder layer for global context + # Uses multi-head self-attention with 8 heads + self.transformer_block = nn.TransformerEncoderLayer( + d_model=out_channels // reduction, + nhead=8 + ) + + # Channel expansion: restore original channel dimension + self.conv2 = nn.Conv2d( + out_channels // reduction, + out_channels, + kernel_size=1, + stride=1, + bias=False + ) + if use_gn: + self.norm2 = nn.GroupNorm(32, out_channels) + else: + self.norm2 = nn.BatchNorm2d(out_channels) + + # Skip connection: adjust channels/spatial dims if needed + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False), + nn.GroupNorm(32, out_channels) if use_gn else nn.BatchNorm2d(out_channels) + ) + + def forward(self, x): + """ + Forward pass through the bottleneck transformer block. + + Process: + 1. Reduce channels with 1x1 convolution + 2. Normalize and activate + 3. Reshape for transformer (flatten spatial dimensions) + 4. Apply self-attention via transformer encoder + 5. Reshape back to spatial format + 6. Expand channels with 1x1 convolution + 7. Add skip connection + 8. Final ReLU activation + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, in_channels, H, W) + + Returns: + torch.Tensor: Output tensor with shape (batch_size, out_channels, H, W) + + Transformer Processing: + The spatial dimensions (H, W) are flattened into a sequence of length H*W, + where each position can attend to all other positions. This allows the + model to capture long-range dependencies that convolutional layers miss. + """ + # Channel reduction + out = F.relu(self.norm1(self.conv1(x))) # (B, C_reduced, H, W) + + # Prepare for transformer: (B, C, H, W) → (H*W, B, C) + b, c, h, w = out.size() + out = out.view(b, c, -1).permute(2, 0, 1) # Flatten spatial dims + + # Self-attention: each position attends to all positions + out = self.transformer_block(out) # (H*W, B, C) + + # Reshape back to spatial format: (H*W, B, C) → (B, C, H, W) + out = out.permute(1, 2, 0).view(b, c, h, w) + + # Channel expansion + out = self.norm2(self.conv2(out)) # (B, out_channels, H, W) + + # Add residual connection + out += self.shortcut(x) + + # Final activation + out = F.relu(out) + return out + +class TransformerDBB(nn.Module): + """ + Dual-Backbone Transformer Network for multi-modal telescope data. + + This architecture uses two independent transformer-augmented backbones to process + different input modalities (charge and timing images) before fusing features for + final prediction. Each backbone combines convolutional layers with transformer + blocks to capture both local patterns and global context. + + Architecture Overview: + Input_1 (charge) → Backbone_1 (Conv + Transformer) → Features_1 + ↓ + Fusion + ↓ + Input_2 (timing) → Backbone_2 (Conv + Transformer) → Features_2 + ↓ + Global Pool → FC → Output + + Each Backbone: + Conv7x7 → Norm → ReLU → + Layer1 (Transformer blocks) → + Layer2 (Transformer blocks, stride=2) → + Layer3 (Transformer blocks, stride=2) → + Layer4 (Transformer blocks, stride=2) → + Adaptive AvgPool + + Attributes: + in_channels (int): Current number of channels (updated during layer construction) + use_gn (bool): Whether to use GroupNorm instead of BatchNorm + use_concat (bool): Whether to concatenate or add backbone outputs + conv1_a, conv1_b (nn.Conv2d): Initial convolutions for each backbone + norm1_a, norm1_b (nn.Module): Initial normalization for each backbone + layer1-4_a, layer1-4_b (nn.Sequential): Transformer block layers for each backbone + dropout (nn.Dropout): Dropout for regularization + adaptive_pool (nn.AdaptiveAvgPool2d): Global average pooling + fc (nn.Linear): Final classification/regression layer + """ + + def __init__(self, block, layers, num_inputs=1, num_classes=1, use_gn=False, use_concat=False, dropout_rate=0.5): + """ + Initialize the dual-backbone transformer network. + + Args: + block: Block class to use (e.g., BottleneckTransformerBlock) + layers (list): Number of blocks in each layer [layer1, layer2, layer3, layer4] + Example: [3, 4, 6, 3] for a ResNet50-like architecture + num_inputs (int, optional): Number of input channels per backbone. Defaults to 1 + num_classes (int, optional): Number of output classes/values. Defaults to 1 + use_gn (bool, optional): Whether to use GroupNorm. Defaults to False + use_concat (bool, optional): Whether to concatenate backbone outputs. + Defaults to False (uses addition) + dropout_rate (float, optional): Dropout probability. Defaults to 0.5 + """ + super(TransformerDBB, self).__init__() + self.in_channels = 64 + self.use_gn = use_gn + self.use_concat = use_concat + + # Backbone 1: Process first input modality (charge images) + self.conv1_a = nn.Conv2d(num_inputs, 64, kernel_size=7, stride=2, padding=3, bias=False) + if use_gn: + self.norm1_a = nn.GroupNorm(32, 64) + else: + self.norm1_a = nn.BatchNorm2d(64) + + # Transformer block layers for backbone 1 + self.layer1_a = self._make_layer(block, 64, layers[0], stride=1) + self.layer2_a = self._make_layer(block, 128, layers[1], stride=2) + self.layer3_a = self._make_layer(block, 256, layers[2], stride=2) + self.layer4_a = self._make_layer(block, 512, layers[3], stride=2) + + # Reset in_channels for second backbone + self.in_channels = 64 + + # Backbone 2: Process second input modality (timing images) + self.conv1_b = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False) + if use_gn: + self.norm1_b = nn.GroupNorm(32, 64) + else: + self.norm1_b = nn.BatchNorm2d(64) + + # Transformer block layers for backbone 2 + self.layer1_b = self._make_layer(block, 64, layers[0], stride=1) + self.layer2_b = self._make_layer(block, 128, layers[1], stride=2) + self.layer3_b = self._make_layer(block, 256, layers[2], stride=2) + self.layer4_b = self._make_layer(block, 512, layers[3], stride=2) + + # Regularization + self.dropout = nn.Dropout(dropout_rate) + + # Global pooling and final layer + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + + # Final FC layer dimension depends on fusion strategy + if self.use_concat: + self.fc = nn.Linear(512 * block.expansion * 2, num_classes) + else: + self.fc = nn.Linear(512 * block.expansion, num_classes) + + def _make_layer(self, block, out_channels, blocks, stride): + """ + Create a layer consisting of multiple transformer blocks. + + Args: + block: Block class (BottleneckTransformerBlock) + out_channels (int): Number of output channels + blocks (int): Number of blocks in this layer + stride (int): Stride for the first block (for downsampling) + + Returns: + nn.Sequential: Sequential container of transformer blocks + """ + layers = [] + # First block may have stride > 1 for downsampling + layers.append(block(self.in_channels, out_channels, stride, use_gn=self.use_gn)) + self.in_channels = out_channels + # Remaining blocks maintain spatial dimensions + for _ in range(1, blocks): + layers.append(block(self.in_channels, out_channels, use_gn=self.use_gn)) + return nn.Sequential(*layers) + + def forward(self, x1): + if x1.shape[1] >= 2: + x1, x2 = torch.split(x1, [1, x1.shape[1]-1], dim=1) + else: + x2 = x1 + """ + Forward pass through the dual-backbone transformer network. + + Process: + 1. Process each input through its backbone (Conv + Transformer layers) + 2. Apply global average pooling to each backbone's output + 3. Flatten spatial dimensions + 4. Apply dropout + 5. Fuse features (concatenate or add) + 6. Final classification/regression layer + + Args: + x1 (torch.Tensor): First input (charge images) + Shape: (batch_size, num_inputs, H, W) + x2 (torch.Tensor): Second input (timing images) + Shape: (batch_size, 1, H, W) + + Returns: + torch.Tensor: Output predictions + Shape: (batch_size, num_classes) + + Feature Fusion: + Concatenation (use_concat=True): + - Preserves all information from both backbones + - Doubles feature dimension + - More parameters in final layer + + Addition (use_concat=False): + - Forces features into shared space + - Fewer parameters + - May lose complementary information + """ + # Backbone 1: Process charge images + out1 = F.relu(self.norm1_a(self.conv1_a(x1))) # Initial conv + out1 = self.layer1_a(out1) # Transformer blocks + out1 = self.layer2_a(out1) + out1 = self.layer3_a(out1) + out1 = self.layer4_a(out1) + out1 = self.adaptive_pool(out1) # Global pooling + out1 = out1.view(out1.size(0), -1) # Flatten + out1 = self.dropout(out1) # Regularization + + # Backbone 2: Process timing images + out2 = F.relu(self.norm1_b(self.conv1_b(x2))) # Initial conv + out2 = self.layer1_b(out2) # Transformer blocks + out2 = self.layer2_b(out2) + out2 = self.layer3_b(out2) + out2 = self.layer4_b(out2) + out2 = self.adaptive_pool(out2) # Global pooling + out2 = out2.view(out2.size(0), -1) # Flatten + out2 = self.dropout(out2) # Regularization + + # Fusion: Combine features from both backbones + if self.use_concat: + out = torch.cat((out1, out2), dim=1) # Concatenate + else: + out = out1 + out2 # Element-wise addition + + # Final prediction + out = self.fc(out) + return out + +def TransformerDuo(num_blocks=[3, 4, 6, 3], num_inputs=1, num_classes=2, use_gn=True, use_concat=False, dropout_rate=0.5): + """ + Factory function to create a Dual-Backbone Transformer network. + + This function instantiates a TransformerDBB model with the specified configuration. + It provides a convenient interface for creating transformer-based models with + different depths and configurations. + + Args: + num_blocks (list, optional): Number of blocks in each layer. Defaults to [3, 4, 6, 3] + Similar to ResNet50 architecture: + - Layer 1: 3 blocks (64 channels) + - Layer 2: 4 blocks (128 channels, stride=2) + - Layer 3: 6 blocks (256 channels, stride=2) + - Layer 4: 3 blocks (512 channels, stride=2) + num_inputs (int, optional): Number of input channels. Defaults to 1 + num_classes (int, optional): Number of output classes/values. Defaults to 2 + use_gn (bool, optional): Whether to use GroupNorm. Defaults to True + Recommended for small batch sizes or distributed training + use_concat (bool, optional): Whether to concatenate backbone outputs. + Defaults to False + dropout_rate (float, optional): Dropout probability. Defaults to 0.5 + + Returns: + TransformerDBB: Instantiated dual-backbone transformer model + + Example: + >>> # Create model for binary classification + >>> model = TransformerDuo( + ... num_blocks=[3, 4, 6, 3], + ... num_inputs=1, + ... num_classes=2, + ... use_gn=True, + ... use_concat=False, + ... dropout_rate=0.5 + ... ) + >>> + >>> # Forward pass + >>> x1 = torch.randn(8, 1, 120, 120) # Charge images + >>> x2 = torch.randn(8, 1, 120, 120) # Timing images + >>> output = model(x1, x2) + >>> print(output.shape) # torch.Size([8, 2]) + + Notes: + - The default configuration creates a lighter model than standard ResNet50 + - Transformer blocks add global context modeling at each layer + - GroupNorm is more stable than BatchNorm for small batches + - Higher dropout rates help prevent overfitting with transformers + """ + return TransformerDBB( + BottleneckTransformerBlock, + num_blocks, + num_inputs, + num_classes=num_classes, + use_gn=use_gn, + use_concat=use_concat, + dropout_rate=dropout_rate + ) \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/Transformers/TransformerDuoSimple.py b/ctlearn/core/pytorch/nets/models/Transformers/TransformerDuoSimple.py new file mode 100644 index 00000000..ac189474 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/Transformers/TransformerDuoSimple.py @@ -0,0 +1,354 @@ +""" +Simplified Dual-Input Transformer Network Module + +This module implements a simplified dual-input transformer architecture that processes +two modalities (charge and timing) by concatenating them at the input level rather than +using separate backbones. This approach is more memory-efficient and has fewer parameters +than the dual-backbone variant. + +The architecture combines convolutional layers with transformer encoder blocks, providing +both local feature extraction and global context modeling through self-attention. + +Classes: + BottleneckTransformerBlock: Bottleneck block with transformer encoder layer + TransformerDBB: Single-backbone transformer for dual-input data + TransformerDuo: Factory function for creating TransformerDBB models + +References: + - "Attention Is All You Need" (Vaswani et al., NeurIPS 2017) + - "BoTNet: Bottleneck Transformers for Visual Recognition" (Srinivas et al., CVPR 2021) +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class BottleneckTransformerBlock(nn.Module): + """ + Bottleneck block with integrated transformer encoder layer. + + This block combines the efficiency of bottleneck architectures with the + global receptive field of transformers. It uses dimensionality reduction + before applying self-attention, making it computationally efficient while + still capturing long-range dependencies. + + Architecture: + Input → Conv1x1(reduce) → Norm → ReLU → + Transformer(self-attention) → + Conv1x1(expand) → Norm → (+) Shortcut → ReLU → Output + + Attributes: + expansion (int): Channel expansion factor (set to 1) + conv1 (nn.Conv2d): 1x1 conv for channel reduction + norm1 (nn.BatchNorm2d or nn.GroupNorm): Normalization after reduction + transformer_block (nn.TransformerEncoderLayer): Self-attention layer + conv2 (nn.Conv2d): 1x1 conv for channel expansion + norm2 (nn.BatchNorm2d or nn.GroupNorm): Normalization after expansion + shortcut (nn.Sequential): Skip connection (identity or 1x1 conv) + """ + + expansion = 1 + + def __init__(self, in_channels, out_channels, stride=1, reduction=4, use_gn=False): + """ + Initialize the bottleneck transformer block. + + Args: + in_channels (int): Number of input channels + out_channels (int): Number of output channels + stride (int, optional): Stride for the shortcut connection. Defaults to 1 + reduction (int, optional): Channel reduction factor. Defaults to 4 + use_gn (bool, optional): Whether to use GroupNorm instead of BatchNorm. + Defaults to False + """ + super(BottleneckTransformerBlock, self).__init__() + + # Channel reduction + self.conv1 = nn.Conv2d( + in_channels, + out_channels // reduction, + kernel_size=1, + stride=1, + bias=False + ) + if use_gn: + self.norm1 = nn.GroupNorm(32, out_channels // reduction) + else: + self.norm1 = nn.BatchNorm2d(out_channels // reduction) + + # Transformer encoder with multi-head self-attention + self.transformer_block = nn.TransformerEncoderLayer( + d_model=out_channels // reduction, + nhead=8 + ) + + # Channel expansion + self.conv2 = nn.Conv2d( + out_channels // reduction, + out_channels, + kernel_size=1, + stride=1, + bias=False + ) + if use_gn: + self.norm2 = nn.GroupNorm(32, out_channels) + else: + self.norm2 = nn.BatchNorm2d(out_channels) + + # Skip connection + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != out_channels: + self.shortcut = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False), + nn.GroupNorm(32, out_channels) if use_gn else nn.BatchNorm2d(out_channels) + ) + + def forward(self, x): + """ + Forward pass through the bottleneck transformer block. + + Args: + x (torch.Tensor): Input tensor with shape (batch_size, in_channels, H, W) + + Returns: + torch.Tensor: Output tensor with shape (batch_size, out_channels, H, W) + """ + # Reduce channels + out = F.relu(self.norm1(self.conv1(x))) + + # Prepare for transformer: (B, C, H, W) → (H*W, B, C) + b, c, h, w = out.size() + out = out.view(b, c, -1).permute(2, 0, 1) + + # Apply self-attention + out = self.transformer_block(out) + + # Reshape back: (H*W, B, C) → (B, C, H, W) + out = out.permute(1, 2, 0).view(b, c, h, w) + + # Expand channels + out = self.norm2(self.conv2(out)) + + # Add residual connection + out += self.shortcut(x) + + # Final activation + out = F.relu(out) + return out + +class TransformerDBB(nn.Module): + """ + Simplified single-backbone transformer for dual-input data. + + This architecture concatenates charge and timing inputs at the channel level + before processing through a single transformer-augmented backbone. This approach + is more parameter-efficient than dual-backbone architectures while still allowing + the model to learn joint representations of both modalities. + + Architecture Overview: + Input_1 (charge) ↘ + → Concatenate → Single Backbone (Conv + Transformer) → Output + Input_2 (timing) ↗ + + Single Backbone: + Concat(x1, x2) → Conv7x7 → Norm → ReLU → + Layer1 (Transformer blocks) → + Layer2 (Transformer blocks, stride=2) → + Layer3 (Transformer blocks, stride=2) → + Layer4 (Transformer blocks, stride=2) → + Adaptive AvgPool → Dropout → FC → Output + + Advantages over Dual-Backbone: + - ~50% fewer parameters (single backbone vs two) + - Lower memory consumption + - Faster training and inference + - Implicit early fusion of modalities + + Attributes: + in_channels (int): Current number of channels (updated during layer construction) + use_gn (bool): Whether to use GroupNorm instead of BatchNorm + conv1 (nn.Conv2d): Initial convolution + norm1 (nn.Module): Initial normalization + layer1-4 (nn.Sequential): Transformer block layers + dropout (nn.Dropout): Dropout for regularization + adaptive_pool (nn.AdaptiveAvgPool2d): Global average pooling + fc (nn.Linear): Final classification/regression layer + """ + + def __init__(self, block, layers, num_inputs=2, num_classes=1, use_gn=False, dropout_rate=0.5): + """ + Initialize the simplified transformer network. + + Args: + block: Block class to use (BottleneckTransformerBlock) + layers (list): Number of blocks in each layer [layer1, layer2, layer3, layer4] + num_inputs (int, optional): Total number of input channels (charge + timing). + Defaults to 2. Should be sum of channels from both modalities. + num_classes (int, optional): Number of output classes/values. Defaults to 1 + use_gn (bool, optional): Whether to use GroupNorm. Defaults to False + dropout_rate (float, optional): Dropout probability. Defaults to 0.5 + + Note: + With num_inputs=2, the network expects concatenated inputs where: + - Channel 0: Charge information + - Channel 1: Timing information + """ + super(TransformerDBB, self).__init__() + self.in_channels = 64 + self.use_gn = use_gn + + # Initial convolution: processes concatenated inputs + # Accepts num_inputs channels (e.g., 2 for charge + timing) + self.conv1 = nn.Conv2d(num_inputs, 64, kernel_size=7, stride=2, padding=3, bias=False) + if use_gn: + self.norm1 = nn.GroupNorm(32, 64) + else: + self.norm1 = nn.BatchNorm2d(64) + + # Transformer block layers + self.layer1 = self._make_layer(block, 64, layers[0], stride=1) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2) + + # Regularization + self.dropout = nn.Dropout(dropout_rate) + + # Global pooling and final layer + self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + def _make_layer(self, block, out_channels, blocks, stride): + """ + Create a layer consisting of multiple transformer blocks. + + Args: + block: Block class (BottleneckTransformerBlock) + out_channels (int): Number of output channels + blocks (int): Number of blocks in this layer + stride (int): Stride for the first block + + Returns: + nn.Sequential: Sequential container of transformer blocks + """ + layers = [] + # First block may downsample + layers.append(block(self.in_channels, out_channels, stride, use_gn=self.use_gn)) + self.in_channels = out_channels + # Remaining blocks maintain dimensions + for _ in range(1, blocks): + layers.append(block(self.in_channels, out_channels, use_gn=self.use_gn)) + return nn.Sequential(*layers) + + def forward(self, x1): + if x1.shape[1] >= 2: + x1, x2 = torch.split(x1, [1, x1.shape[1]-1], dim=1) + else: + x2 = x1 + """ + Forward pass through the simplified transformer network. + + Process: + 1. Concatenate inputs along channel dimension + 2. Process through single backbone (Conv + Transformer layers) + 3. Global average pooling + 4. Dropout for regularization + 5. Final classification/regression layer + + Args: + x1 (torch.Tensor): First input (charge images) + Shape: (batch_size, 1, H, W) + x2 (torch.Tensor): Second input (timing images) + Shape: (batch_size, 1, H, W) + + Returns: + torch.Tensor: Output predictions + Shape: (batch_size, num_classes) + + Input Fusion: + Early fusion via concatenation: x = [x1 | x2] + This allows the network to learn joint features from both modalities + from the very first layer, unlike dual-backbone approaches where + fusion happens later. + + Example: + >>> model = TransformerDBB(...) + >>> x1 = torch.randn(16, 1, 120, 120) # Charge + >>> x2 = torch.randn(16, 1, 120, 120) # Timing + >>> output = model(x1, x2) + >>> print(output.shape) # torch.Size([16, 1]) + """ + # Concatenate inputs along channel dimension + # x1: (B, 1, H, W), x2: (B, 1, H, W) → x: (B, 2, H, W) + x = torch.cat((x1, x2), dim=1) + + # Process through single backbone + out = F.relu(self.norm1(self.conv1(x))) # Initial conv + out = self.layer1(out) # Transformer blocks + out = self.layer2(out) + out = self.layer3(out) + out = self.layer4(out) + + # Global pooling and prediction + out = self.adaptive_pool(out) # (B, 512, 1, 1) + out = out.view(out.size(0), -1) # (B, 512) + out = self.dropout(out) # Regularization + out = self.fc(out) # (B, num_classes) + + return out + +def TransformerDuo(num_blocks=[3, 4, 6, 3], num_inputs=2, num_classes=2, use_gn=True, dropout_rate=0.3): + """ + Factory function to create a simplified dual-input transformer network. + + This function instantiates a TransformerDBB model with the specified configuration. + It provides a convenient interface for creating transformer-based models that + process two input modalities through early fusion. + + Args: + num_blocks (list, optional): Number of blocks in each layer. Defaults to [3, 4, 6, 3] + Similar to ResNet50 architecture + num_inputs (int, optional): Total input channels. Defaults to 2 + Should equal the sum of channels from all input modalities + num_classes (int, optional): Number of output classes/values. Defaults to 2 + use_gn (bool, optional): Whether to use GroupNorm. Defaults to True + Recommended for stability with small batches + dropout_rate (float, optional): Dropout probability. Defaults to 0.3 + Lower than dual-backbone default (0.5) since single backbone + + Returns: + TransformerDBB: Instantiated simplified transformer model + + Example: + >>> # Create model for binary classification + >>> model = TransformerDuo( + ... num_blocks=[3, 4, 6, 3], + ... num_inputs=2, + ... num_classes=2, + ... use_gn=True, + ... dropout_rate=0.3 + ... ) + >>> + >>> # Forward pass + >>> x1 = torch.randn(8, 1, 120, 120) # Charge + >>> x2 = torch.randn(8, 1, 120, 120) # Timing + >>> output = model(x1, x2) + >>> print(output.shape) # torch.Size([8, 2]) + + Comparison with Dual-Backbone: + Simplified (this): + - Pros: Fewer parameters, faster, lower memory + - Cons: Less modality-specific feature learning + + Dual-Backbone: + - Pros: More modality-specific features, potentially higher accuracy + - Cons: More parameters, slower, higher memory + """ + return TransformerDBB( + BottleneckTransformerBlock, + num_blocks, + num_inputs=num_inputs, + num_classes=num_classes, + use_gn=use_gn, + dropout_rate=dropout_rate + ) \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/__init__.py b/ctlearn/core/pytorch/nets/models/__init__.py new file mode 100644 index 00000000..d4860865 --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/__init__.py @@ -0,0 +1,18 @@ + +import ctlearn.core.pytorch.nets.models.ThinResNet_DBB.ThinResNet_DBB +import ctlearn.core.pytorch.nets.models.ThinResNet.ThinResNet +import ctlearn.core.pytorch.nets.models.DBBRegNet.DBBRegNet +import ctlearn.core.pytorch.nets.models.DoubleBBEfficientNet.DoubleBBEfficientNet + +import ctlearn.core.pytorch.nets.models.DBBDanet.DBBDanet + +# Diffusion +import ctlearn.core.pytorch.nets.models.NoPropDT.NoPropDT +import ctlearn.core.pytorch.nets.models.NoPropDTReg.NoPropDTReg +import ctlearn.core.pytorch.nets.models.NoPropDTRegDBB.NoPropDTRegDBB +import ctlearn.core.pytorch.nets.models.DBBNoPropDTReg.DBBNoPropDTReg +import ctlearn.core.pytorch.nets.models.NoPropDTReg2.NoPropDTReg2 + +# Hourglass +import ctlearn.core.pytorch.nets.models.StackedHGNet.StackedHGNet +import ctlearn.core.pytorch.nets.models.StackedHGNetDBB.StackedHGNetDBB \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/models/gcn/gcn.py b/ctlearn/core/pytorch/nets/models/gcn/gcn.py new file mode 100644 index 00000000..45fd4c9d --- /dev/null +++ b/ctlearn/core/pytorch/nets/models/gcn/gcn.py @@ -0,0 +1,206 @@ +""" +Graph Convolutional Network (GCN) Module + +This module implements a Graph Convolutional Network for processing graph-structured +data from Cherenkov telescope observations. The GCN architecture is particularly useful +for analyzing the spatial relationships between triggered pixels in telescope cameras. + +The network uses graph convolutions to propagate information between neighboring pixels, +followed by graph pooling and classification/regression heads. + +Classes: + GCN: Graph Convolutional Network for telescope data analysis + +References: + - "Semi-Supervised Classification with Graph Convolutional Networks" (Kipf & Welling, ICLR 2017) + - PyTorch Geometric: https://pytorch-geometric.readthedocs.io/ +""" + +from torch_geometric.loader import DataLoader +from torch_geometric.nn import GATConv +import pickle +import torch +import torch.nn.functional as F +from torch_geometric.nn import GCNConv, GraphConv +from torch_geometric.nn import global_mean_pool, global_max_pool +from torch.nn import Linear, Softmax, PReLU + +class GCN(torch.nn.Module): + """ + Graph Convolutional Network for telescope camera data. + + This network processes graph-structured data where nodes represent camera pixels + and edges connect neighboring pixels. The architecture consists of: + 1. Multiple graph convolutional layers to aggregate information + 2. Global pooling to obtain graph-level representation + 3. Fully connected layers for final prediction + + Graph Structure: + - Nodes: Camera pixels with features (e.g., charge, timing) + - Edges: Connections between neighboring pixels + - Graph: Complete telescope camera image + + Architecture: + Input Graph → GraphConv1 → PReLU → Dropout → + GraphConv2 → PReLU → Dropout → GraphConv3 → + Global Mean Pool → Linear → Linear → Output + + Attributes: + conv0 (GCNConv): Initial graph convolution (currently unused) + conv1 (GraphConv): First graph convolutional layer + conv2 (GraphConv): Second graph convolutional layer + conv3 (GraphConv): Third graph convolutional layer + lin_0 (Linear): First fully connected layer + lin_1 (Linear): Second fully connected layer (output) + prelu_1 (PReLU): Parametric ReLU activation for conv1 + prelu_2 (PReLU): Parametric ReLU activation for conv2 + """ + + def __init__(self, hidden_channels, num_node_features=1, num_outputs=1): + """ + Initialize the Graph Convolutional Network. + + Args: + hidden_channels (int): Number of hidden features in graph conv layers + Typical values: 64, 128, 256 + Higher values allow more complex representations but increase computation + + num_node_features (int, optional): Number of features per node (pixel). + Defaults to 1 (e.g., charge only) + Can be 2 for charge + timing, or more for additional features + + num_outputs (int, optional): Number of output values. Defaults to 1 + For classification: 2 (gamma vs proton) + For regression: 1 (energy or direction component) + + Network Design Choices: + - GraphConv vs GCNConv: GraphConv is more general, supports edge features + - Bias enabled: Helps with different graph sizes and structures + - 3 conv layers: Balances receptive field vs oversmoothing + - PReLU: Learnable activation that can adapt to data + - Dropout (p=0.3): Regularization to prevent overfitting + + Example: + >>> # Create GCN for binary classification + >>> model = GCN(hidden_channels=128, num_node_features=2, num_outputs=2) + >>> + >>> # Create GCN for energy regression + >>> model = GCN(hidden_channels=64, num_node_features=1, num_outputs=1) + """ + super(GCN, self).__init__() + + # Set random seed for reproducibility + torch.manual_seed(12345) + + # Graph convolutional layers + use_bias = True # Enable bias for better expressiveness + + # Initial convolution (currently unused, kept for potential skip connections) + self.conv0 = GCNConv(num_node_features, hidden_channels, bias=use_bias) + + # First graph convolution: node features → hidden_channels + self.conv1 = GraphConv(num_node_features, hidden_channels, bias=use_bias) + + # Second graph convolution: hidden_channels → hidden_channels + # Aggregates information from 2-hop neighbors + self.conv2 = GraphConv(hidden_channels, hidden_channels, bias=use_bias) + + # Third graph convolution: hidden_channels → hidden_channels + # Aggregates information from 3-hop neighbors + self.conv3 = GraphConv(hidden_channels, hidden_channels, bias=use_bias) + + # Fully connected layers for prediction + # First FC: Processes pooled graph representation + self.lin_0 = Linear(hidden_channels, hidden_channels) + + # Output FC: Maps to final predictions + self.lin_1 = Linear(hidden_channels, num_outputs) + + # Parametric ReLU activations (learnable negative slope) + self.prelu_1 = PReLU() # After conv1 + self.prelu_2 = PReLU() # After conv2 + + def forward(self, x, edge_index, batch): + """ + Forward pass through the Graph Convolutional Network. + + Process: + 1. Apply graph convolutions to propagate information between neighbors + 2. Use PReLU activation and dropout after each conv layer + 3. Perform global pooling to get graph-level representation + 4. Apply fully connected layers for final prediction + + Args: + x (torch.Tensor): Node feature matrix with shape (num_nodes, num_node_features) + Each row represents features of one pixel in the camera + Example: For 1000 triggered pixels with 2 features each: (1000, 2) + + edge_index (torch.Tensor): Graph connectivity in COO format + Shape: (2, num_edges) + edge_index[0]: Source nodes + edge_index[1]: Target nodes + Example: [[0, 1, 1], [1, 0, 2]] represents edges 0→1, 1→0, 1→2 + + batch (torch.Tensor): Batch vector which assigns each node to a graph + Shape: (num_nodes,) + Example: [0, 0, 0, 1, 1, 2] for 3 graphs with 3, 2, and 1 nodes + + Returns: + torch.Tensor: Output predictions with shape (batch_size, num_outputs) + For classification: logits before softmax + For regression: predicted values + + Graph Convolution Process: + Each conv layer aggregates information from neighbors: + h_i^(l+1) = σ(Σ_{j∈N(i)} (h_j^(l) · W^(l)) / √(d_i · d_j)) + where: + - h_i: features of node i + - N(i): neighbors of node i + - W: learnable weight matrix + - d_i: degree of node i + - σ: activation function (PReLU) + + Example: + >>> # Process a batch of 3 graphs + >>> x = torch.randn(100, 2) # 100 total pixels, 2 features each + >>> edge_index = torch.randint(0, 100, (2, 300)) # 300 edges + >>> batch = torch.tensor([0]*30 + [1]*40 + [2]*30) # 3 graphs + >>> + >>> output = model(x, edge_index, batch) + >>> print(output.shape) # torch.Size([3, 1]) for 3 graphs, 1 output each + """ + # 1. Obtain node embeddings through graph convolutions + + # First convolution: Extract local patterns + x = self.conv1(x, edge_index) # (num_nodes, hidden_channels) + x = self.prelu_1(x) # Parametric ReLU activation + x = F.dropout(x, p=0.3, training=self.training) # Regularization + + # Second convolution: Aggregate 2-hop neighborhood information + x = self.conv2(x, edge_index) # (num_nodes, hidden_channels) + x = self.prelu_2(x) # Parametric ReLU activation + x = F.dropout(x, p=0.3, training=self.training) # Regularization + + # Third convolution: Aggregate 3-hop neighborhood information + x = self.conv3(x, edge_index) # (num_nodes, hidden_channels) + + # Note: Skip connections commented out + # x = x + x_ori # Residual connection with initial features + # Could help with gradient flow and preserve initial information + + # 2. Readout layer: Aggregate node features to graph-level representation + # Global mean pooling: Average all node features per graph + x = global_mean_pool(x, batch=batch) # (batch_size, hidden_channels) + + # Alternative: Global max pooling (commented out) + # x = global_max_pool(x, batch=batch) + # Max pooling captures most salient features but loses information + + # 3. Apply fully connected layers for final prediction + # First FC layer: Further process graph representation + x = self.lin_0(x) # (batch_size, hidden_channels) + + # Output layer: Map to final predictions + x = self.lin_1(x) # (batch_size, num_outputs) + + return x \ No newline at end of file diff --git a/ctlearn/core/pytorch/nets/optimizer/__init__.py b/ctlearn/core/pytorch/nets/optimizer/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/pytorch/nets/optimizer/optimizer.py b/ctlearn/core/pytorch/nets/optimizer/optimizer.py new file mode 100644 index 00000000..7bafd015 --- /dev/null +++ b/ctlearn/core/pytorch/nets/optimizer/optimizer.py @@ -0,0 +1,129 @@ +""" +Learning Rate Scheduling Functions Module + +This module provides learning rate scheduling functions for training neural networks. +It includes implementations of various learning rate schedules commonly used in deep +learning, particularly focusing on one-cycle and cosine annealing schedules that have +been shown to improve training convergence and final model performance. + +Functions: + one_cycle: Generate a one-cycle learning rate schedule with sinusoidal ramp + +References: + - "A disciplined approach to neural network hyper-parameters" (Smith, 2018) + - "Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates" (Smith & Topin, 2017) +""" + +import math + +def one_cycle(y1=0.0, y2=1.0, steps=100): + """ + Generate a one-cycle learning rate schedule using a sinusoidal ramp. + + This function creates a lambda function that implements a smooth, cosine-based + transition from an initial learning rate (y1) to a maximum learning rate (y2) + over a specified number of steps. The schedule follows a half-cosine curve, + providing a gradual warm-up and smooth transition. + + The one-cycle policy has been shown to enable faster training and better + generalization by allowing the use of higher learning rates during training + while maintaining stability through the smooth schedule. + + Mathematical Formula: + lr(x) = ((1 - cos(x * π / steps)) / 2) * (y2 - y1) + y1 + + where: + - x: current step (0 to steps) + - lr(x): learning rate at step x + - The cosine creates a smooth S-shaped curve from y1 to y2 + + Args: + y1 (float, optional): Initial learning rate (start value). Defaults to 0.0 + Typically set to a small value or 0 for warm-up from zero + y2 (float, optional): Maximum learning rate (end value). Defaults to 1.0 + This is the peak learning rate to reach during training + Should be set based on learning rate range tests + steps (int, optional): Total number of steps for the schedule. Defaults to 100 + This determines how quickly the learning rate increases + For epoch-based: steps = num_epochs + For iteration-based: steps = num_epochs * batches_per_epoch + + Returns: + function: A lambda function that takes step index x and returns the + corresponding learning rate. The function signature is: + lambda x: float + + Schedule Characteristics: + - At x=0: Returns y1 (starting learning rate) + - At x=steps/2: Returns (y1+y2)/2 (midpoint) + - At x=steps: Returns y2 (maximum learning rate) + - Smooth acceleration: No sudden jumps in learning rate + - Cosine-based: Provides gentle start and smooth transition + + Usage Patterns: + 1. One-Cycle Policy (Smith, 2018): + - Phase 1: Ramp up from low LR to high LR (this function) + - Phase 2: Ramp down from high LR to very low LR (inverse) + + 2. Warm-up Only: + - Use this function alone for gradual LR warm-up + - Helps stabilize training at the beginning + + Example: + >>> # Create a schedule from 0.0 to 0.1 over 1000 steps + >>> schedule = one_cycle(y1=0.0, y2=0.1, steps=1000) + >>> + >>> # Get learning rate at step 0 (start) + >>> lr_start = schedule(0) + >>> print(f"LR at step 0: {lr_start:.6f}") # 0.000000 + >>> + >>> # Get learning rate at step 500 (halfway) + >>> lr_mid = schedule(500) + >>> print(f"LR at step 500: {lr_mid:.6f}") # ~0.050000 + >>> + >>> # Get learning rate at step 1000 (end) + >>> lr_end = schedule(1000) + >>> print(f"LR at step 1000: {lr_end:.6f}") # 0.100000 + + >>> # Use with PyTorch LambdaLR scheduler + >>> import torch.optim as optim + >>> optimizer = optim.SGD(model.parameters(), lr=1.0) + >>> lambda_func = one_cycle(y1=0.0, y2=1.0, steps=100) + >>> scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_func) + >>> + >>> # During training + >>> for epoch in range(100): + ... train_one_epoch() + ... scheduler.step() # Updates LR according to schedule + + >>> # Complete one-cycle schedule (warm-up + cool-down) + >>> warmup = one_cycle(0.0, 1.0, steps=50) + >>> cooldown = one_cycle(1.0, 0.01, steps=50) + >>> # Use warmup for first 50 epochs, cooldown for last 50 + + Notes: + - The returned lambda is stateless; it only depends on the input step + - Can be used with PyTorch's LambdaLR scheduler for automatic LR updates + - Works with both epoch-based and iteration-based schedules + - The cosine curve provides smoother transitions than linear schedules + - Avoids sudden LR changes that can destabilize training + + Common Configurations: + Warm-up from zero: + >>> schedule = one_cycle(0.0, 0.001, steps=10) # 10-epoch warm-up + + One-cycle for 100 epochs: + >>> up = one_cycle(0.0, 0.1, steps=70) # Ramp up: epochs 0-70 + >>> down = one_cycle(0.1, 0.0001, steps=30) # Ramp down: epochs 70-100 + + Fine-tuning with low LR: + >>> schedule = one_cycle(0.00001, 0.0001, steps=20) # Gentle increase + + See Also: + torch.optim.lr_scheduler.LambdaLR: PyTorch scheduler using lambda functions + torch.optim.lr_scheduler.OneCycleLR: Built-in one-cycle implementation + torch.optim.lr_scheduler.CosineAnnealingLR: Cosine annealing schedule + """ + # Return lambda function that computes LR for any step x + # The cosine function creates a smooth S-curve from y1 to y2 + return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1 diff --git a/ctlearn/core/pytorch/utils/__init__.py b/ctlearn/core/pytorch/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/core/pytorch/utils/utils.py b/ctlearn/core/pytorch/utils/utils.py new file mode 100644 index 00000000..04a246bd --- /dev/null +++ b/ctlearn/core/pytorch/utils/utils.py @@ -0,0 +1,1050 @@ +import importlib +import logging +import os +# import pkg_resources +import sys +import time +import pickle +import numpy as np +import pandas as pd +import tables +import yaml +import zlib +import math +import re +import matplotlib.pyplot as plt +import torch +from astropy.coordinates import SkyCoord, AltAz +import astropy.units as u +try: + from ctapipe_io_lst.constants import LST1_LOCATION +except ImportError: + from astropy.coordinates import EarthLocation + # Fallback coordinates for LST-1 + LST1_LOCATION = EarthLocation(lon=-17.89139 * u.deg, lat=28.76139 * u.deg, height=2184 * u.m) + +from ctlearn.core.ctlearn_enum import Task, Mode +import time as time_ +from ctapipe.io import read_table, write_table +from astropy.table import Table +crab_nebula_config = { + "observation_mode": "wobble", + "n_off_wobble": 1, + "source_name": "Crab Nebula", + "source_ra": 83.63308333, + "source_dec": 22.0145 +} +from astropy.time import Time +from ctapipe.coordinates import CameraFrame +try: + # Python 3.8+ + from importlib.metadata import version as get_version +except ImportError: + # For Python <3.8, install backport: importlib-metadata + from importlib_metadata import version as get_version + + +def clip_alt(alt): + """ + Make sure altitude is not larger than 90 deg (it happens in some MC files for zenith=0), + to keep astropy happy + """ + return np.clip(alt, -90.0 * u.deg, 90.0 * u.deg) + +def radec_to_camera(sky_coordinate, obstime, pointing_alt, pointing_az, focal): + """ + Coordinate transform from sky coordinate to camera coordinates (x, y) in distance + + Parameters + ---------- + sky_coordinate: astropy.coordinates.sky_coordinate.SkyCoord + obstime: astropy.time.Time + pointing_alt: pointing altitude in angle unit + pointing_az: pointing altitude in angle unit + focal: astropy Quantity + + Returns + ------- + camera frame: `astropy.coordinates.sky_coordinate.SkyCoord` + """ + + horizon_frame = AltAz(location=LST1_LOCATION, obstime=obstime) + + pointing_direction = SkyCoord( + alt=clip_alt(pointing_alt), az=pointing_az, frame=horizon_frame,unit="deg" + ) + + camera_frame = CameraFrame( + focal_length=focal, + telescope_pointing=pointing_direction, + obstime=obstime, + location=LST1_LOCATION, + ) + + camera_pos = sky_coordinate.transform_to(camera_frame) + + return camera_pos + +def get_expected_source_pos(data, data_type, config, effective_focal_length=29.30565 * u.m): + + # For real data + if data_type == 'real_data': + # source is always at the ceter of camera for ON mode + if config.get('observation_mode') == 'on': + expected_src_pos_x_m = np.zeros(len(data)) + expected_src_pos_y_m = np.zeros(len(data)) + + # compute source position in camera coordinate event by event for wobble mode + elif config.get('observation_mode') == 'wobble': + + if 'source_name' in config: + source_coord = SkyCoord.from_name(config.get('source_name')) + elif 'source_ra' and 'source_dec' in config: + source_coord = SkyCoord(config.get('source_ra'), config.get('source_dec'), frame="icrs", unit="deg") + else: + raise KeyError( + 'source position (`source_name` or `source_ra` & `source_dec`) is not defined in a config file for source-dependent analysis.' + ) + + time = data['dragon_time'] + obstime = Time(time, scale='utc', format='unix') + # pointing_alt = u.Quantity(data['alt_tel'], u.rad, copy=False) + # pointing_az = u.Quantity(data['az_tel'], u.rad, copy=False) + pointing_alt_deg = data['pointing_alt']*u.deg + pointing_az_deg = data['pointing_az']*u.deg + + pointing_alt_rad = pointing_alt_deg.to(u.rad) + pointing_az_rad = pointing_az_deg.to(u.rad) + + source_pos = radec_to_camera(source_coord, obstime, pointing_alt_deg, pointing_az_deg, effective_focal_length) + + expected_src_pos_x_m = source_pos.x.to_value(u.m) + expected_src_pos_y_m = source_pos.y.to_value(u.m) + + else: + raise KeyError( + '`observation_mode` is not defined in a config file for source-dependent analysis. It should be `on` or `wobble`' + ) + + return expected_src_pos_x_m, expected_src_pos_y_m + +def sky_to_camera(alt, az, focal, pointing_alt, pointing_az, horizon_frame): + """ + Coordinate transform from aky position (alt, az) (in angles) + to camera coordinates (x, y) in distance. + + Parameters + ---------- + alt: astropy Quantity + az: astropy Quantity + focal: astropy Quantity + pointing_alt: pointing altitude in angle unit + pointing_az: pointing altitude in angle unit + + Returns + ------- + camera frame: `astropy.coordinates.sky_coordinate.SkyCoord` + """ + pointing_direction = SkyCoord( + alt=clip_alt(pointing_alt), az=pointing_az, frame=horizon_frame + ) + + camera_frame = CameraFrame( + focal_length=focal, telescope_pointing=pointing_direction + ) + + event_direction = SkyCoord(alt=clip_alt(alt), az=az, frame=horizon_frame) + + camera_pos = event_direction.transform_to(camera_frame) + + return camera_pos +#------------------------------------------------------------------------------------------------------------------- +def reco_src_sky_to_camera(data, effective_focal_length=29.30565 * u.m): + + time = data['dragon_time'] + obstime = Time(time, scale='utc', format='unix') + # horizon_frame = AltAz(location=LST1_LOCATION, obstime=obstime) + + + # time = data['utc_time'] + # time_utc = Time(time, format="mjd", scale="tai") + # obstime = time_utc.utc + + alt_deg = u.Quantity(data['alt'], u.deg, copy=False) + az_deg = u.Quantity(data['az'], u.deg, copy=False) + + horizon_frame = AltAz(location=LST1_LOCATION, obstime=obstime) + + tel_alt_deg = u.Quantity(data['tel_pointing_alt'], u.deg, copy=False) + tel_az_deg = u.Quantity(data['tel_pointing_az'], u.deg, copy=False) + + source_pos_in_camera = sky_to_camera( + alt_deg, + az_deg, + effective_focal_length, + tel_alt_deg, + tel_az_deg, + horizon_frame=horizon_frame, + ) + expected_src_pos_x_m = source_pos_in_camera.x.to_value(u.m) + expected_src_pos_y_m = source_pos_in_camera.y.to_value(u.m) + + # ----------------------------------------------------------------- + # TODO: Remove this. It is the equivalent to the code above + # tel_alt_rad = tel_alt_deg.to(u.rad) + # tel_az_rad = tel_az_deg.to(u.rad) + # time = data['utc_time'] + # time_utc = Time(time, format="mjd", scale="tai") + # time_utc = time_utc.utc + # telescope_pointing = SkyCoord(alt=tel_alt_rad, az=tel_az_rad, + # frame=AltAz(obstime=time_utc, + # location=LST1_LOCATION)) + + # source_pos = SkyCoord( + # alt=clip_alt(alt_deg), az=az_deg, frame=horizon_frame + # ) + + # # CameraFrame is terribly slow without the erfa interpolator below... + # # with erfa_astrom.set(ErfaAstromInterpolator(5 * u.min)): + # camera_frame = CameraFrame(focal_length=effective_focal_length, + # telescope_pointing=telescope_pointing, + # location=LST1_LOCATION, obstime=time_utc) + + # source_pos_camera = source_pos.transform_to(camera_frame) + + return expected_src_pos_x_m, expected_src_pos_y_m +#------------------------------------------------------------------------------------------------------------------- +expected_structure = { + "data": { + "train_gamma_proton": None, + "validation_gamma_proton": None, + "train_gamma": None, + "validation_gamma": None, + "test_gamma": None, + "test_proton": None, + "test_electron": None, + "test_validation_gamma": None, + "test_validation_gamma_proton": None + }, + "run_details": { + "mode": None, + "task": None, + "test_type": None, + "experiment_number": None + }, + "cut-off": { + "leakage_intensity": None, + "intensity": None + }, + "model": { + "model_type": { + "model_name": None, + "parameters": None, + }, + "model_energy": { + "model_name": None, + "parameters": None, + }, + "model_direction": { + "model_name": None, + "parameters": None, + } + }, + "hyp": { + "epochs": None, + "batches": None, + "dynamic_batches": None, + "optimizer": None, + "momentum": None, + "weight_decay": None, + "learning_rate": None, + "lrf": None, + "start_epoch": None, + "steps_epoch": None, + "l2_lambda": None, + "adam_epsilon": None, + "gradient_clip_val": None, + "save_k": None + }, + "augmentation": { + "use_augmentation": None, + "aug_prob": None, + "rot_prob": None, + "trans_prob": None, + "flip_hor_prob": None, + "flip_ver_prob": None, + "mask_prob": None, + "noise_prob": None, + "max_rot": None, + "max_trans": None + }, + "normalization": { + "use_clean": None, + "type_mu": None, + "type_sigma": None, + "dir_mu": None, + "dir_sigma": None, + "energy_mu": None, + "energy_sigma": None, + "apply_log_scaling": [None, None] + }, + "dataset": { + "num_workers": None, + "pin_memory": None, + "persistent_workers": None + }, + "arch": { + "device": None, + "precision_type": None, + "precision_energy": None, + "precision_direction": None, + "devices": None, + "strategy": None + }, +} +#------------------------------------------------------------------------------------------------------------------- +# Sanity check function +def sanity_check(config, expected_structure): + """ + Recursively checks if all required keys in the expected_structure exist in the config data. + Raises a KeyError if a key is missing. + """ + for key, substructure in expected_structure.items(): + if key not in config: + raise KeyError(f"Missing key: {key}") + + # If the substructure is a dictionary, recursively check the subkeys + if isinstance(substructure, dict): + if not isinstance(config[key], dict): + raise KeyError(f"Expected a dictionary for key: {key}, but got: {type(config[key])}") + sanity_check(config[key], substructure) + +#------------------------------------------------------------------------------------------------------------------- +def extract_loss_value(file_path): + # Use regular expression to find the pattern + match = re.search(r'_([0-9]+\.[0-9]+)\.pth$', file_path) + if match: + return match.group(1) # Return the whole match + else: + return None +#------------------------------------------------------------------------------------------------------------------- +def convert_ascii_list_to_string(ascii_values, padding_value=0): + # Filter out padding values and convert each ASCII value to its corresponding character + characters = [chr(ascii_val) + for ascii_val in ascii_values if ascii_val != padding_value] + + # Join all characters to form the string + return ''.join(characters) +#------------------------------------------------------------------------------------------------------------------- +def extract_accuracy_from_filename(filename): + # Use a regular expression to find numbers that might be in the format of floating point numbers + match = re.search(r"(\d+\.\d+)", filename) + if match: + # Return the first occurrence of a floating point number as a float + return float(match.group(1)) + else: + # If no matching number format is found, handle it appropriately + return None +#------------------------------------------------------------------------------------------------------------------- +def find_bin_index(value, bins): + index = np.digitize([value], bins) - 1 + index = max(0, min(index[0], len(bins) - 1)) # Ensure the index is within the valid range + bin_value = bins[index] + return index, bin_value +#------------------------------------------------------------------------------------------------------------------- +def get_bin_value(index:int, bins): + # Ensure the index is within the valid range + index = max(0, min(index, len(bins) - 1)) + # Return the bin value associated with the index + bin_value = bins[index] + return bin_value +#------------------------------------------------------------------------------------------------------------------- +def get_bin_value(indices:np.array, bins): + # Ensure indices are within the valid range + indices = np.clip(indices, 0, len(bins) - 1) + # Return the bin values associated with the indices + return np.array(bins)[indices] +#------------------------------------------------------------------------------------------------------------------- +def get_bin_value(indices:torch.tensor, bins): + # Convert bins to a tensor if they aren't already + # bins_tensor = torch.tensor(bins) + # Ensure indices are within the valid range + indices = torch.clamp(indices, 0, len(bins) - 1) + + # Return the bin values associated with the indices + return bins[indices] + +#------------------------------------------------------------------------------------------------------------------- +def decompress_data(compressed_data, data_shape, data_type=np.float16): + + decompressed_data = zlib.decompress(compressed_data) + restored_array = np.frombuffer(decompressed_data, dtype=data_type) + restored_array_reshaped = np.reshape(restored_array, data_shape) + + return restored_array_reshaped +#------------------------------------------------------------------------------------------------------------------- +def compress_data(data): + # Convert to bytes + bytes_data = data.tobytes() + + # Array Compression + compressed_data = zlib.compress(bytes_data) + return compressed_data +#------------------------------------------------------------------------------------------------------------------- +def create_experiment_folder(prefix="run_", next_number=None): + + + """ + Create the next folder within the specified directory with a given prefix. + If next_number is not specified, automatically find the next available number. + + :param run_directory: The directory where folders are managed. + :param prefix: Prefix used for folders. + :param next_number: Optional. Specify the number to be used for the new folder. + """ + run_directory="./run" + + # Ensure the 'run' directory exists + if not os.path.exists(run_directory): + os.makedirs(run_directory) + print(f"Directory '{run_directory}' created.") + + if next_number is None: + # List all subdirectories in the 'run' directory + folders = [f for f in os.listdir(run_directory) if os.path.isdir(os.path.join(run_directory, f))] + # Filter folders that match the prefix pattern and end with a digit + matching_folders = [f for f in folders if f.startswith(prefix) and f[len(prefix):].isdigit()] + + # Find the highest number and calculate the next one + if matching_folders: + highest_number = max(int(folder[len(prefix):]) for folder in matching_folders) + next_number = highest_number + 1 + else: + next_number = 0 + + # Create the new folder with the next number + new_folder_name = f"{prefix}{next_number}" + new_folder_path = os.path.join(run_directory, new_folder_name) + new_folder_path+="/" + if not os.path.exists(new_folder_path): + os.makedirs(new_folder_path) + print(f"New folder created: {new_folder_path}") + return new_folder_path +#------------------------------------------------------------------------------------------------------------ +def cartesian_to_alt_az(cartesian): + x = cartesian[0] + y = cartesian[1] + z = cartesian[2] + # Calculate the radius + r = math.sqrt(x**2 + y**2 + z**2) + + # Calculate altitude in radians + if r == 0: # Avoid division by zero + altitude_rad = 0 + else: + altitude_rad = math.asin(z / r) + + # Calculate azimuth in radians + azimuth_rad = math.atan2(y, x) + + # Adjust azimuth to be within the range [0, 2*pi) + # if azimuth_rad < 0: + # azimuth_rad += 2 * math.pi + + return np.array([altitude_rad, azimuth_rad, r]) +#------------------------------------------------------------------------------------------------------------ +def alt_az_to_cartesian(altitude_rad, azimuth_rad, r=1): + + # Estimate Cartesian coordinates + x = r * math.cos(altitude_rad) * math.cos(azimuth_rad) + y = r * math.cos(altitude_rad) * math.sin(azimuth_rad) + z = r * math.sin(altitude_rad) + + return np.array([x, y, z]) +#------------------------------------------------------------------------------------------------------------ +def load_pickle(pickle_file): + with open(pickle_file, 'rb') as file: + print(f"Loading pickle file :{pickle_file}") + return pickle.load(file) + +#------------------------------------------------------------------------------------------------------------ +def create_key_value_array(data_dict, id): + """ + Creates an array of key-value pairs from a dictionary where each key has a list of values. + + Args: + data_dict (dict): The dictionary containing keys with lists of values. + + Returns: + list: An array of key-value pairs, where each pair is a tuple containing the key and one value. + """ + + key_value_array = [] + for key, value_list in data_dict.items(): + # Handle cases where a key has an empty list: + if value_list.size==0: + continue # Skip keys with empty lists + + # Choose the first value from the list: + value = value_list[id] + + # Create a tuple (key, value) and append it to the array + key_value_array.append(value) + + return key_value_array + +#------------------------------------------------------------------------------------------------------------ +def decompress_data(compressed_data, data_shape, data_type=np.float16): + + decompressed_data = zlib.decompress(compressed_data) + restored_array = np.frombuffer(decompressed_data, dtype=data_type) + restored_array_reshaped = np.reshape(restored_array, data_shape) + + return restored_array_reshaped +#------------------------------------------------------------------------------------------------------------ +# def setup_logging(config, log_dir, debug, log_to_file): + +# # Log configuration to a text file in the log dir +# time_str = time.strftime("%Y%m%d_%H%M%S") +# config_filename = os.path.join(log_dir, time_str + "_config.yml") +# with open(config_filename, "w") as outfile: +# ctlearn_version = pkg_resources.get_distribution("ctlearn").version +# tensorflow_version = pkg_resources.get_distribution("tensorflow").version +# outfile.write( +# "# Training performed with " +# "CTLearn version {} and TensorFlow version {}.\n".format( +# ctlearn_version, tensorflow_version +# ) +# ) +# yaml.dump(config, outfile, default_flow_style=False) + +# # Set up logger +# logger = logging.getLogger() + +# if debug: +# logger.setLevel(logging.DEBUG) +# else: +# logger.setLevel(logging.INFO) + +# logger.handlers = [] # remove existing handlers from any previous runs +# if not log_to_file: +# handler = logging.StreamHandler() +# else: +# logging_filename = os.path.join(log_dir, time_str + "_logfile.log") +# handler = logging.FileHandler(logging_filename) +# handler.setFormatter(logging.Formatter("%(levelname)s:%(message)s")) +# logger.addHandler(handler) + +# return logger +def setup_logging(config, log_dir, debug, log_to_file): + # Log configuration to a text file in the log dir + time_str = time.strftime("%Y%m%d_%H%M%S") + config_filename = os.path.join(log_dir, time_str + "_config.yml") + with open(config_filename, "w") as outfile: + try: + ctlearn_version = get_version("ctlearn") + except Exception: + ctlearn_version = "unknown" + try: + tensorflow_version = get_version("tensorflow") + except Exception: + tensorflow_version = "unknown" + outfile.write( + "# Training performed with " + "CTLearn version {} and TensorFlow version {}.\n".format( + ctlearn_version, tensorflow_version + ) + ) + yaml.dump(config, outfile, default_flow_style=False) + + # Set up logger + logger = logging.getLogger() + + if debug: + logger.setLevel(logging.DEBUG) + else: + logger.setLevel(logging.INFO) + + logger.handlers = [] # remove existing handlers from any previous runs + if not log_to_file: + handler = logging.StreamHandler() + else: + logging_filename = os.path.join(log_dir, time_str + "_logfile.log") + handler = logging.FileHandler(logging_filename) + handler.setFormatter(logging.Formatter("%(levelname)s:%(message)s")) + logger.addHandler(handler) + + return logger +#------------------------------------------------------------------------------------------------------------ +def setup_DL1DataReader(config, mode): + + # Parse file list or prediction file list + if mode in ["train", "load_only"]: + if isinstance(config["Data"]["file_list"], str): + data_files = [] + with open(config["Data"]["file_list"]) as f: + for line in f: + line = line.strip() + if line and line[0] != "#": + data_files.append(line) + config["Data"]["file_list"] = data_files + if not isinstance(config["Data"]["file_list"], list): + raise ValueError( + "Invalid file list '{}'. " + "Must be list or path to file".format(config["Data"]["file_list"]) + ) + else: + file_list = config["Prediction"]["prediction_file_lists"][ + config["Prediction"]["prediction_file"] + ] + if file_list.endswith(".txt"): + data_files = [] + with open(file_list) as f: + for line in f: + line = line.strip() + if line and line[0] != "#": + data_files.append(line) + config["Data"]["file_list"] = data_files + elif file_list.endswith(".h5"): + config["Data"]["file_list"] = [file_list] + if not isinstance(config["Data"]["file_list"], list): + raise ValueError( + "Invalid prediction file list '{}'. " + "Must be list or path to file".format(file_list) + ) + + mc_file = True + with tables.open_file(config["Data"]["file_list"][0], mode="r") as f: + if "CTA PRODUCT DATA MODEL NAME" in f.root._v_attrs: + data_format = "stage1" + elif "dl1_data_handler_version" in f.root._v_attrs: + data_format = "dl1dh" + else: + raise ValueError( + "Data format is not implemented in the DL1DH reader. Available data formats are 'stage1' and 'dl1dh'." + ) + if data_format == "dl1dh" and "source_name" in f.root._v_attrs: + mc_file = False + + allow_overwrite = config["Data"].get("allow_overwrite", True) + if "allow_overwrite" in config["Data"]: + del config["Data"]["allow_overwrite"] + + selected_telescope_types = config["Data"]["selected_telescope_types"] + camera_types = [tel_type.split("_")[-1] for tel_type in selected_telescope_types] + + tasks = config["Reco"] + transformations = [] + event_info = [] + if data_format == "dl1dh": + if "parameter_list" not in config["Data"] and mode == "predict": + config["Data"]["parameter_list"] = [ + "hillas_intensity", + "log_hillas_intensity", + "hillas_x", + "hillas_y", + "hillas_r", + "hillas_phi", + "hillas_length", + "hillas_width", + "hillas_psi", + "hillas_skewness", + "leakage_intensity_width_1", + "leakage_intensity_width_2", + "morphology_num_islands", + "impact", + "log_impact", + "maxheight", + "log_maxheight", + "cherenkovdensity", + "log_hillasintensity_over_cherenkovdensity", + "cherenkovradius", + "impact_over_cherenkovradius", + "p1grad", + "sqrt_p1grad_p1grad", + ] + # Parse list of event selection filters + event_selection = {} + for s in config["Data"].get("event_selection", {}): + s = {"module": "dl1_data_handler.filters", **s} + filter_fn, filter_params = load_from_module(**s) + event_selection[filter_fn] = filter_params + config["Data"]["event_selection"] = event_selection + + # Parse list of image selection filters + image_selection = {} + for s in config["Data"].get("image_selection", {}): + s = {"module": "dl1_data_handler.filters", **s} + filter_fn, filter_params = load_from_module(**s) + image_selection[filter_fn] = filter_params + config["Data"]["image_selection"] = image_selection + + if "direction" in tasks: + event_info.append("src_pos_cam_x") + event_info.append("src_pos_cam_y") + transformations.append( + { + "name": "AltAz", + "args": { + "alt_col_name": "src_pos_cam_x", + "az_col_name": "src_pos_cam_y", + "deg2rad": False, + }, + } + ) + else: + if "parameter_list" not in config["Data"] and mode == "predict": + config["Data"]["parameter_list"] = [ + "hillas_intensity", + "hillas_fov_lon", + "hillas_fov_lat", + "hillas_r", + "hillas_phi", + "hillas_length", + "hillas_length_uncertainty", + "hillas_width", + "hillas_width_uncertainty", + "hillas_psi", + "hillas_skewness", + "hillas_kurtosis", + "timing_slope", + "timing_intercept", + "timing_deviation", + "leakage_pixels_width_1", + "leakage_pixels_width_2", + "leakage_intensity_width_1", + "leakage_intensity_width_2", + "concentration_cog", + "concentration_core", + "concentration_pixel", + "morphology_n_pixels", + "morphology_n_islands", + "morphology_n_medium_islands", + "morphology_n_large_islands", + "intensity_max", + "intensity_min", + "intensity_mean", + "intensity_std", + "intensity_skewness", + "intensity_kurtosis", + "peak_time_max", + "peak_time_min", + "peak_time_mean", + "peak_time_std", + "peak_time_skewness", + "peak_time_kurtosis", + ] + if "direction" in tasks: + event_info.append("true_alt") + event_info.append("true_az") + transformations.append({"name": "DeltaAltAz_fix_subarray"}) + + if "particletype" in tasks: + event_info.append("true_shower_primary_id") + + if "energy" in tasks: + if mc_file: + event_info.append("true_energy") + transformations.append({"name": "MCEnergy"}) + + concat_telescopes = config["Input"].get("concat_telescopes", False) + if config["Data"]["mode"] == "stereo" and not concat_telescopes: + for tel_desc in selected_telescope_types: + transformations.append( + { + "name": "SortTelescopes", + "args": {"sorting": "size", "tel_desc": f"{tel_desc}"}, + } + ) + + # Convert interpolation image shapes from lists to tuples, if present + if "interpolation_image_shape" in config["Data"].get("mapping_settings", {}): + config["Data"]["mapping_settings"]["interpolation_image_shape"] = { + k: tuple(l) + for k, l in config["Data"]["mapping_settings"][ + "interpolation_image_shape" + ].items() + } + + if allow_overwrite: + config["Data"]["event_info"] = event_info + config["Data"]["mapping_settings"]["camera_types"] = camera_types + else: + transformations = config["Data"].get("transforms", {}) + + transforms = [] + # Parse list of Transforms + for t in transformations: + t = {"module": "dl1_data_handler.transforms", **t} + transform, args = load_from_module(**t) + transforms.append(transform(**args)) + config["Data"]["transforms"] = transforms + + # Possibly add additional info to load if predicting to write later + if mode == "predict": + + if "Prediction" not in config: + config["Prediction"] = {} + if "event_info" not in config["Data"]: + config["Data"]["event_info"] = [] + config["Data"]["event_info"].extend(["event_id", "obs_id"]) + if data_format == "dl1dh" and not mc_file: + config["Data"]["event_info"].extend(["mjd", "milli_sec", "nano_sec"]) + + return config["Data"], data_format +#------------------------------------------------------------------------------------------------------------ +def load_from_module(name, module, path=None, args=None): + if path is not None and path not in sys.path: + sys.path.append(path) + mod = importlib.import_module(module) + fn = getattr(mod, name) + params = args if args is not None else {} + return fn, params +#------------------------------------------------------------------------------------------------------------ +def recover_alt_az(fix_pointing,alt_off, az_off): + + az_off_deg = u.Quantity(az_off, unit=u.rad).to(u.deg).value + alt_off_deg = u.Quantity(alt_off, unit=u.rad).to(u.deg).value + + reco_direction = fix_pointing.spherical_offsets_by( + [az_off_deg] * u.deg, + [alt_off_deg] * u.deg + ) + + # Clamp altitude offset to avoid exceeding valid range + # alt_off_deg = np.clip(alt_off_deg, -90, 90) + + return reco_direction +#------------------------------------------------------------------------------------------------------------ +def write_output(h5file, data, predictions, labels, task:Task, mode:Mode): + prediction_dir = h5file.replace(f'{h5file.split("/")[-1]}', "") + if not os.path.exists(prediction_dir): + os.makedirs(prediction_dir) + + # Store dl2 data + reco = {} + if os.path.isfile(h5file): + with pd.HDFStore(h5file, mode="r") as file: + h5file_keys = list(file.keys()) + if f"/dl2/reco" in h5file_keys: + reco = pd.read_hdf(file, key=f"/dl2/reco") + + + reco["event_id"] = np.array(data["event_id"]) + reco["obs_id"] = np.array(data["obs_id"]) + + + + if task ==Task.type: + + for n, name in enumerate(data["class_names"]): + reco[name + "ness"] = np.array(predictions["type"][:, n]) + # reco["type_feature_vector"] = np.array(predictions["type_feature_vector"]) + reco["reco_type"]= np.array(predictions["type_class"])*101 + + if mode != Mode.observation:#"observation": + if data["energy_unit"] == "log(TeV)": + reco["true_energy"] = np.power(10, labels["true_energy"]) + else: + reco["true_energy"] = labels["true_energy"] + + + if task == Task.energy: + # if data["energy_unit"] == "log(TeV)" or np.min(predictions["energy"]) < 0.0: + # reco["reco_energy"] = np.power(10, predictions["energy"][:, 0]) + # reco["log_reco_energy"] =predictions["energy"][:, 0] + # else: + reco["reco_energy"] = np.power(10, predictions["energy"][:, 0]) + reco["log_reco_energy"] =predictions["energy"][:, 0] + + + if mode != Mode.observation:#"observation": + reco["true_alt"] = np.float32(np.rad2deg(labels["true_alt_az"][:,0])) + reco["true_az"] = np.float32(np.rad2deg(labels["true_alt_az"][:,1])) + + if task==Task.direction: + # reco["reco_alt"] = np.array(predictions["direction"][:, 1]) + # reco["reco_az"] = np.array(predictions["direction"][:, 0]) + + reco_az, reco_alt = [], [] + reco_az_off, reco_alt_off = [], [] + reco_src_x, reco_src_y = [], [] + dragon_time = [] + utc_time = [] + + data_type = 'real_data' + if mode == Mode.observation: + # reco["reco_src_x"] = [] + # reco["reco_src_y"] = [] + reco["src_x"] = data["pointing"]["src_x"] + reco["src_y"] = data["pointing"]["src_y"] + + reco_data = {} + if mode != Mode.observation:#"observation": + pointing_alt = np.full(len(data["obs_id"]), data["pointing"]["pointing_alt"]) + pointing_az = np.full(len(data["obs_id"]), data["pointing"]["pointing_az"]) + + reco["alt_tel"] = pointing_alt + reco["az_tel"] = pointing_az + else: + pointing_alt = data["pointing"]["pointing_alt"] + pointing_az = data["pointing"]["pointing_az"] + + reco["alt_tel"] = pointing_alt + reco["az_tel"] = pointing_az + + # fix_pointing_t = SkyCoord( + # pointing_az * u.deg, + # pointing_alt * u.deg, + # frame="altaz", + # unit="deg", + # ) + + horizon_frame="altaz" + + if mode== Mode.observation: + time = data["pointing"]["dragon_time"] + obstime = Time(time, scale='utc', format='unix') + horizon_frame = AltAz(location=LST1_LOCATION, obstime=obstime) + + # obstime_lst = Time("2018-11-01T02:00") + # horizon_frame_lst = AltAz(location=LST1_LOCATION, obstime=obstime_lst) + + # fix_pointing_lst = SkyCoord( + # az=pointing_az * u.deg, + # alt=pointing_alt * u.deg, + # frame=horizon_frame_lst, + # unit="deg", + # ) + + fix_pointing = SkyCoord( + az=pointing_az * u.deg, + alt=pointing_alt * u.deg, + frame=horizon_frame, + unit="deg", + ) + # ---------------------------------------------------------------------------------------------------------------------------------------- + # start = time_.time() + sigma_az=[] + sigma_alt=[] + if len(predictions["direction"])>0: + az_off_ = predictions["direction"][:, 0] + alt_off_ = predictions["direction"][:, 1] + + else: + az_off_ = predictions["direction_mu"][:, 0] + alt_off_ = predictions["direction_mu"][:, 1] + sigma_az = predictions["direction_sigma"][:, 0] + sigma_alt = predictions["direction_sigma"][:, 1] + + reco_direction_ = recover_alt_az(fix_pointing,alt_off_,az_off_) + reco_az_item_=reco_direction_.az.to_value(u.deg)[0] + reco_alt_item_=reco_direction_.alt.to_value(u.deg)[0] + + reco_az=reco_az_item_ + reco_alt = reco_alt_item_ + reco_az_off = az_off_ + reco_alt_off = alt_off_ + + reco_data_ = {} + if mode == Mode.observation: + reco_data_["alt"] = reco_alt + reco_data_["az"] = reco_az + reco_data_["tel_pointing_alt"] = fix_pointing.alt.value + reco_data_["tel_pointing_az"] = fix_pointing.az.value + + reco_data_["dragon_time"] = data["pointing"]["dragon_time"] + reco_data_["utc_time"] = data["pointing"]["utc_time"] + + reco_src_ = reco_src_sky_to_camera(reco_data_,effective_focal_length=data["effective_focal_length"]) + + reco_src_x = reco_src_[0] + reco_src_y = reco_src_[1] + dragon_time = data["pointing"]["dragon_time"] + # New fix + # dragon_time = Time(dragon_time, format='unix_tai') + + utc_time = data["pointing"]["utc_time"] + + # end = time_.time() + # print("Optimized: ", end - start) + # ---------------------------------------------------------------------------------------------------------------------------------------- + if mode == Mode.observation: + reco["reco_src_x"] = reco_src_x + reco["reco_src_y"] = reco_src_y + reco["dragon_time"] = dragon_time + reco["utc_time"] = utc_time + + reco["reco_alt"] = reco_alt + reco["reco_az"] = reco_az + reco["reco_alt_off"] = reco_alt_off + reco["reco_az_off"] = reco_az_off + reco["sigma_alt"]= sigma_alt + reco["sigma_az"]= sigma_az + + del predictions + import gc + gc.collect() + + # Convertir reco a Astropy Table + reco_table = Table() + for key, val in reco.items(): + reco_table[key] = np.array(val) + + if data["include_nsb_patches"] is None: + pd.DataFrame(data=reco).to_hdf(h5file, key=f"/dl2/reco", mode="a", format="table") + # write_table( + # reco_table, + # h5file, + # f"/dl2/reco", + # overwrite=True, + # ) + else: + # write_table( + # reco_table, + # h5file, + # f"/trigger/reco", + # overwrite=True, + # ) + pd.DataFrame(data=reco).to_hdf(h5file, key=f"/trigger/reco", mode="a", format="table") + + + + # Guardar usando ctapipe.io.write_table + + + + # Store the simulation information for pyirf + if data["simulation_info"] and data["include_nsb_patches"] != "all": + pd.DataFrame(data=data["simulation_info"] , index=[0]).to_hdf( + h5file, key=f"/info/mc_header", mode="a", format="table") + + + # Store the selected Hillas parameters (dl1b) + if data["parameter_names"] and data["include_nsb_patches"] != "all": + tel_counter = 0 + if data["mode"]== "mono": + tel_type = list(data["selected_telescopes"].keys())[0] + tel_ids = "tel" + for tel_id in data["selected_telescopes"][tel_type]: + tel_ids += f"_{tel_id}" + parameters = {} + for p, parameter in enumerate(data["parameter_names"]): + parameter_list = np.array(data["parameter_data"])[:,p] + + parameters[parameter] = parameter_list + pd.DataFrame(data=parameters).to_hdf( + h5file, key=f"/dl1b/{tel_type}/{tel_ids}", mode="a", format="table") + else: + for tel_type in data["selected_telescopes"]: + for t, tel_id in enumerate(data["selected_telescopes"][tel_type]): + parameters = {} + for p, parameter in enumerate(data["parameter_list"]): + parameter_list = np.array(data["parameter_list"])[:, tel_counter + t, p] + + parameters[parameter] = parameter_list + pd.DataFrame(data=parameters).to_hdf( + h5file, key=f"/dl1b/{tel_type}/tel_{tel_id}", mode="a", format="table") + tel_counter += len(data["selected_telescopes"][tel_type]) + + print(f"File saved: {h5file}") \ No newline at end of file diff --git a/ctlearn/core/pytorch/utils/utils_torch.py b/ctlearn/core/pytorch/utils/utils_torch.py new file mode 100644 index 00000000..1613798c --- /dev/null +++ b/ctlearn/core/pytorch/utils/utils_torch.py @@ -0,0 +1,300 @@ +""" +PyTorch Utility Functions Module + +This module provides utility functions for coordinate transformations, optimizer +management, and model weight tracking in PyTorch. These utilities are specifically +designed for Cherenkov telescope array analysis in CTLearn. + +Functions: + cartesian_to_alt_az: Convert Cartesian coordinates to altitude/azimuth + alt_az_to_cartesian: Convert altitude/azimuth to Cartesian coordinates + adjust_learning_rate: Modify learning rate for all optimizer parameter groups + compare_weights: Compare current model weights with initial weights +""" + +import torch + +def cartesian_to_alt_az(directions): + """ + Convert 3D Cartesian direction vectors to altitude and azimuth angles. + + This function transforms Cartesian coordinates (x, y, z) representing + unit direction vectors into spherical coordinates (altitude, azimuth) + commonly used in astronomical coordinate systems. + + Coordinate System: + - Cartesian: (x, y, z) where x² + y² + z² = r² + - Spherical: (altitude, azimuth) in radians + * Altitude (elevation): angle above the horizontal plane [-π/2, π/2] + * Azimuth: angle in the horizontal plane, measured from x-axis [−π, π] + + Args: + directions (torch.Tensor): Batch of 3D direction vectors + Shape: (batch_size, 3) where each row is [x, y, z] + Can be unit vectors or any length (will be normalized internally) + + Returns: + torch.Tensor: Altitude and azimuth angles in radians + Shape: (batch_size, 2) where each row is [altitude, azimuth] + - altitude: Range [-π/2, π/2] radians (-90° to 90°) + - azimuth: Range [-π, π] radians (-180° to 180°) + + Mathematical Formulation: + r = √(x² + y² + z²) + altitude = arcsin(z / r) + azimuth = arctan2(y, x) + + Example: + >>> # Pointing straight up (zenith) + >>> directions = torch.tensor([[0.0, 0.0, 1.0]]) + >>> alt_az = cartesian_to_alt_az(directions) + >>> print(alt_az) # [[π/2, 0]] + + >>> # Batch of directions + >>> directions = torch.tensor([ + ... [1.0, 0.0, 0.0], # East horizon + ... [0.0, 1.0, 0.0], # North horizon + ... [0.0, 0.0, 1.0], # Zenith + ... ]) + >>> alt_az = cartesian_to_alt_az(directions) + >>> # Result: [[0, 0], [0, π/2], [π/2, 0]] + + Notes: + - Handles zero vectors safely by replacing r=0 with r=1 to avoid division by zero + - Azimuth follows standard convention: 0 at x-axis, increases counterclockwise + - For shower direction reconstruction in gamma-ray astronomy + - Compatible with astropy coordinate transformations + """ + # Calculate magnitude (radius) of each direction vector + r = torch.sqrt(torch.sum(directions**2, dim=1)) + + # Prevent division by zero for null vectors + # Replace zero magnitudes with 1.0 to avoid NaN in division + safe_r = torch.where(r == 0, torch.tensor(1.0, device=r.device), r) + + # Calculate altitude (elevation angle above horizontal plane) + # altitude = arcsin(z / r), range: [-π/2, π/2] + altitude_rad = torch.asin(directions[:, 2] / safe_r) + + # Calculate azimuth (angle in horizontal plane from x-axis) + # azimuth = arctan2(y, x), range: [-π, π] + azimuth_rad = torch.atan2(directions[:, 1], directions[:, 0]) + + # Note: Azimuth normalization to [0, 2π] is commented out + # Current implementation returns azimuth in [-π, π] + # Uncomment the following line if [0, 2π] range is needed: + # azimuth_rad = torch.where(azimuth_rad < 0, azimuth_rad + 2 * torch.pi, azimuth_rad) + + # Stack altitude and azimuth into single tensor + # Shape: (batch_size, 2) + return torch.stack((altitude_rad, azimuth_rad), dim=1) + +def alt_az_to_cartesian(altitude_rad, azimuth_rad, r=1): + """ + Convert altitude and azimuth angles to 3D Cartesian coordinates. + + This function transforms spherical coordinates (altitude, azimuth) into + Cartesian coordinates (x, y, z). This is the inverse operation of + cartesian_to_alt_az. + + Coordinate System: + - Input: (altitude, azimuth, radius) in radians and distance units + - Output: (x, y, z) Cartesian coordinates + + Args: + altitude_rad (torch.Tensor): Altitude angles in radians + Shape: (batch_size,) or scalar + Range: [-π/2, π/2] (where 0 is horizon, π/2 is zenith) + + azimuth_rad (torch.Tensor): Azimuth angles in radians + Shape: (batch_size,) or scalar + Range: [-π, π] or [0, 2π] + Convention: 0 at x-axis, increases counterclockwise + + r (float or torch.Tensor, optional): Radial distance from origin + Default: 1 (unit vectors) + Can be scalar (same for all) or tensor (batch_size,) + + Returns: + torch.Tensor: 3D Cartesian direction vectors + Shape: (batch_size, 3) where each row is [x, y, z] + If r=1, returns unit vectors + + Mathematical Formulation: + x = r · cos(altitude) · cos(azimuth) + y = r · cos(altitude) · sin(azimuth) + z = r · sin(altitude) + + Example: + >>> # Convert zenith direction (straight up) + >>> altitude = torch.tensor([np.pi/2]) # 90 degrees + >>> azimuth = torch.tensor([0.0]) + >>> xyz = alt_az_to_cartesian(altitude, azimuth) + >>> print(xyz) # [[0, 0, 1]] + + >>> # Convert horizon direction pointing East + >>> altitude = torch.tensor([0.0]) + >>> azimuth = torch.tensor([0.0]) + >>> xyz = alt_az_to_cartesian(altitude, azimuth) + >>> print(xyz) # [[1, 0, 0]] + + >>> # Batch conversion with custom radius + >>> altitudes = torch.tensor([0.0, np.pi/4, np.pi/2]) + >>> azimuths = torch.tensor([0.0, np.pi/2, 0.0]) + >>> xyz = alt_az_to_cartesian(altitudes, azimuths, r=2.0) + >>> # Result: [[2, 0, 0], [0, √2, √2], [0, 0, 2]] + + Notes: + - Default r=1 produces unit vectors + - Useful for converting predicted angles back to 3D directions + - Compatible with shower direction reconstruction in CTA analysis + - Inverse of cartesian_to_alt_az (within numerical precision) + """ + # Calculate x-coordinate + # x = r · cos(altitude) · cos(azimuth) + # Points in the horizontal plane, aligned with azimuth angle + x = r * torch.cos(altitude_rad) * torch.cos(azimuth_rad) + + # Calculate y-coordinate + # y = r · cos(altitude) · sin(azimuth) + # Points in the horizontal plane, perpendicular to x + y = r * torch.cos(altitude_rad) * torch.sin(azimuth_rad) + + # Calculate z-coordinate + # z = r · sin(altitude) + # Points vertically (altitude component) + z = r * torch.sin(altitude_rad) + + # Stack coordinates into 3D vectors + # Shape: (batch_size, 3) + return torch.stack((x, y, z), dim=1) + +def adjust_learning_rate(optimizer, lr): + """ + Adjust the learning rate for all parameter groups in an optimizer. + + This function modifies the learning rate of all parameter groups in a + PyTorch optimizer. Useful for implementing custom learning rate schedules, + warm-up strategies, or manual learning rate adjustments during training. + + Args: + optimizer (torch.optim.Optimizer): PyTorch optimizer instance + Can be any optimizer (SGD, Adam, AdamW, etc.) + Contains one or more parameter groups + + lr (float): New learning rate to set + Must be positive + Applied to all parameter groups uniformly + + Side Effects: + Modifies the 'lr' field of all parameter groups in the optimizer in-place + + Example: + >>> # Initialize optimizer + >>> optimizer = torch.optim.Adam(model.parameters(), lr=0.001) + >>> + >>> # Reduce learning rate by factor of 10 + >>> adjust_learning_rate(optimizer, 0.0001) + >>> + >>> # Verify new learning rate + >>> print(optimizer.param_groups[0]['lr']) + 0.0001 + + >>> # Use in training loop for manual scheduling + >>> for epoch in range(num_epochs): + ... if epoch == 50: + ... adjust_learning_rate(optimizer, lr * 0.1) + ... train_one_epoch() + + Notes: + - Affects ALL parameter groups (if optimizer has multiple groups) + - For different learning rates per group, access param_groups directly + - Common use cases: + * Learning rate warm-up + * Manual learning rate decay + * Cyclical learning rate schedules + * Recovery from divergence during training + - Alternative: Use PyTorch lr_scheduler classes for automatic scheduling + + See Also: + torch.optim.lr_scheduler: Built-in learning rate schedulers + """ + # Iterate over all parameter groups in the optimizer + for param_group in optimizer.param_groups: + # Update the learning rate for this parameter group + param_group['lr'] = lr + +def compare_weights(model, initial_weights): + """ + Compare current model weights with initial weights to detect changes. + + This utility function checks which parameters in a PyTorch model have + changed compared to their initial values. Useful for debugging training + issues, verifying that training is updating weights, or checking if + certain layers are frozen correctly. + + Args: + model (torch.nn.Module): PyTorch model to check + Can be any neural network model + + initial_weights (dict): Dictionary of initial weight tensors + Format: {parameter_name: tensor} + Typically obtained from model.state_dict() before training + Example: initial_weights = {name: param.data.clone() + for name, param in model.named_parameters()} + + Side Effects: + Prints the names of parameters that have changed + Does not modify the model or weights + + Example: + >>> # Save initial weights before training + >>> model = ResNet(num_outputs=2) + >>> initial_weights = { + ... name: param.data.clone() + ... for name, param in model.named_parameters() + ... } + >>> + >>> # Train for one epoch + >>> train_one_epoch(model, optimizer, train_loader) + >>> + >>> # Check which weights changed + >>> compare_weights(model, initial_weights) + Weight changed: conv1.weight + Weight changed: bn1.weight + Weight changed: layer1.0.conv1.weight + ... + + >>> # Check if frozen layers stayed frozen + >>> for name, param in model.named_parameters(): + ... if 'backbone' in name: + ... param.requires_grad = False + >>> + >>> train_one_epoch(model, optimizer, train_loader) + >>> compare_weights(model, initial_weights) + # Should not print any 'backbone' layers + + Notes: + - Only checks parameters that exist in both model and initial_weights + - Uses torch.equal() for exact comparison (no tolerance) + - Useful for debugging: + * Verifying training is working (weights should change) + * Checking frozen layers (weights should NOT change) + * Identifying which layers are being updated + * Detecting gradient flow issues + - For large models, consider checking only specific layers + + Performance: + - Comparison is done in-place, no additional memory allocation + - Fast for most models, but can be slow for very large models + """ + # Iterate over all named parameters in the model + for name, param in model.named_parameters(): + # Get the corresponding initial weight tensor + initial_weight = initial_weights[name] + + # Compare current weight with initial weight + # torch.equal() returns True only if tensors are exactly equal + if not torch.equal(initial_weight, param.data): + # Print parameter name if it has changed + print(f"Weight changed: {name}") \ No newline at end of file diff --git a/ctlearn/core/pytorch/visualization/tsne_test.py b/ctlearn/core/pytorch/visualization/tsne_test.py new file mode 100644 index 00000000..1f75f3c5 --- /dev/null +++ b/ctlearn/core/pytorch/visualization/tsne_test.py @@ -0,0 +1,28 @@ +import numpy as np +import matplotlib.pyplot as plt +from sklearn import datasets +from sklearn.manifold import TSNE + +if __name__ == "__main__": + # Load the Iris dataset + iris = datasets.load_iris() + X = iris.data + y = iris.target + + # Create a TSNE instance with desired parameters + tsne = TSNE(n_components=2, random_state=42) + + # Perform TSNE + X_embedded = tsne.fit_transform(X) + + # Plot the embedded points with matplotlib + plt.figure(figsize=(8, 6)) + colors = ['red', 'blue', 'green'] + for i, color in enumerate(colors): + plt.scatter(X_embedded[y == i, 0], X_embedded[y == i, 1], c=color, label=iris.target_names[i]) + + plt.legend(loc='best') + plt.title('TSNE visualization of the Iris dataset') + plt.xlabel('TSNE 1') + plt.ylabel('TSNE 2') + plt.show() \ No newline at end of file diff --git a/ctlearn/core/pytorch/visualization/vis_utils.py b/ctlearn/core/pytorch/visualization/vis_utils.py new file mode 100644 index 00000000..9e1b0c80 --- /dev/null +++ b/ctlearn/core/pytorch/visualization/vis_utils.py @@ -0,0 +1,195 @@ +import numpy as np +from sklearn.metrics import roc_curve, auc +import matplotlib.pyplot as plt +import cv2 +from io import BytesIO +import os +try: + import seaborn as sns +except ImportError: + sns = None +import pandas as pd +import astropy.units as u +try: + import ctaplot +except ImportError: + ctaplot = None +import torch +# ---------------------------------------------------------------------------------------------------------- +def plot_energy_resolution_error(val_energy_pred_list,val_energy_label_list,val_hillas_intensity_list): + # Create a DataFrame to store the data + data = pd.DataFrame({ + 'energy_reco': list(val_energy_pred_list), + 'energy_true': list(val_energy_label_list), + 'hillas_intensity': list(val_hillas_intensity_list) + }) + if ctaplot is None: + raise ImportError("ctaplot is required for plot_energy_resolution_error. Install it via 'pip install ctaplot'") + + cut_off = 50 + filtered_data = data[data['hillas_intensity'] > cut_off] + true_energy = u.Quantity(filtered_data['energy_true'], u.TeV) + reco_energy = u.Quantity(filtered_data['energy_reco'], u.TeV) + # .apply(lambda x: x[0]) + + fig, ax = plt.subplots() + ctaplot.plot_energy_resolution(true_energy, reco_energy, label="Energy resolution", ax=ax) + # Skip warning + ax.get_xaxis().set_units(None) + ax.get_yaxis().set_units(None) + + ctaplot.plot_energy_resolution_cta_requirement('north', ax=ax, color='black') + ax.legend() + ax.set_ylim(bottom=0, top=1.5) + # plt.show(block=True) + # plt.close() + + return fig +# ---------------------------------------------------------------------------------------------------------- +def plot_direction_resolution_error(val_alt_pred_list,val_az_pred_list, val_alt_label_list,val_az_label_list, val_energy_label_list,val_hillas_intensity_list): + + # Create a DataFrame to store the data + data = pd.DataFrame({ + 'alt_reco': list(val_alt_pred_list), + 'az_reco': list(val_az_pred_list), + 'alt_true': list(val_alt_label_list), + 'az_true': list(val_az_label_list), + 'energy_true': list(val_energy_label_list), + 'hillas_intensity': list(val_hillas_intensity_list) + }) + if ctaplot is None: + raise ImportError("ctaplot is required for plot_direction_resolution_error. Install it via 'pip install ctaplot'") + + cut_off = 50 + filtered_data = data[data['hillas_intensity'] > cut_off] + true_energy = u.Quantity(filtered_data['energy_true'], u.TeV) + + alt_true = u.Quantity(filtered_data['alt_true'], u.rad) + az_true = u.Quantity(filtered_data['az_true'], u.rad) + + alt_reco = u.Quantity(filtered_data['alt_reco'], u.rad) + az_reco = u.Quantity(filtered_data['az_reco'], u.rad) + + fig, ax = plt.subplots() + ax = ctaplot.plot_angular_resolution_per_energy(alt_true, alt_reco, az_true, az_reco, true_energy, label="Ang res (mono)") + ctaplot.plot_angular_resolution_cta_requirement('north', ax=ax, color='black') + ax.legend() + ax.set_ylim(bottom=0, top=1.5) + return fig +# ---------------------------------------------------------------------------------------------------------- +def plot_confusion_matrix(cm, classes, cm_file_name, save_folder): + """ + Plots and saves the confusion matrix. + + Args: + cm (torch.Tensor or np.ndarray): The confusion matrix. + classes (list): List of class names. + cm_file_name (str): Name of the file to save. + save_folder (str): Folder to save the plot. + """ + if isinstance(cm, torch.Tensor): + cm = cm.cpu().numpy() # Convert to NumPy if it's a tensor + + if sns is None: + raise ImportError("seaborn is required for plot_confusion_matrix. Install it via 'pip install seaborn'") + + plt.figure(figsize=(10, 7)) + sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', + xticklabels=classes, yticklabels=classes) + plt.xlabel('Predicted Labels') + plt.ylabel('True Labels') + + # Compute class-wise accuracies + class_accuracies = cm.diagonal() / cm.sum(axis=1) + accuracy_str = "Accuracy: \n" + for idx, class_name in enumerate(classes): + accuracy_str += f" {class_name}: {round(class_accuracies[idx] * 100, 2)}%\n" + + plt.title('Confusion Matrix\n' + accuracy_str) + + # Ensure the save folder exists + os.makedirs(save_folder, exist_ok=True) + + # Save the figure + plt.savefig(os.path.join(save_folder, f"{cm_file_name}.png")) + plt.close() + return class_accuracies +# ---------------------------------------------------------------------------------------------------------- +def create_image_mosaic(images_list, text_list, image_id, rows=4, cols=4, save_path='./',save_name='mosaic'): + """ + Creates an image mosaic using Matplotlib, saves it, and returns it as a numpy array. + + Args: + images_list (list): List of numpy array images. + rows (int): Number of rows in the mosaic. + cols (int): Number of columns in the mosaic. + save_path (str): Path to save the mosaic image. + + Returns: + numpy.ndarray: Image of the mosaic loaded as a numpy array for OpenCV. + """ + fig, axes = plt.subplots(rows, cols, figsize=( + 12, 12)) # Adjust the figure size as needed + axes = axes.ravel() + + for idx, ax in enumerate(axes): + if idx < len(images_list): + ax.set_title(text_list[idx], fontsize=6) + ax.imshow(images_list[idx], cmap='viridis' if len( + images_list[idx].shape) == 2 else None) + ax.axis('off') + else: + ax.axis('off') + + plt.tight_layout() + + # Save the figure to a buffer + buf = BytesIO() + plt.savefig(buf, format='png') + buf.seek(0) + image = np.frombuffer(buf.getvalue(), dtype=np.uint8) + buf.close() + + # Close the plot to free memory + plt.close(fig) + + # Convert buffer to OpenCV image format + image = cv2.imdecode(image, cv2.IMREAD_COLOR) + + # Optionaly save the image to disk + if save_name and save_path: + cv2.imwrite(os.path.join(save_path,save_name+"_"+str(image_id)+".png"), image) + + return image +# ---------------------------------------------------------------------------------------------------------- +def plot_roc_and_calculate_auc(ground_truth, predicted_probabilities): + """ + This function calculates the ROC curve and the AUC given a vector of true labels and predicted probabilities. + + Parameters: + - ground_truth: numpy array with true labels (0s and 1s). + - predicted_probabilities: numpy array with the probabilities of the positive class. + + Returns: + - AUC score. + """ + # Calcular puntos para la curva ROC + fpr, tpr, _ = roc_curve(ground_truth, predicted_probabilities) + # Calcular el AUC + roc_auc = auc(fpr, tpr) + + # Graficar la curva ROC + plt.figure() + plt.plot(fpr, tpr, color='darkorange', + lw=2, label='ROC curve (area = %0.2f)' % roc_auc) + plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--') + plt.xlim([0.0, 1.0]) + plt.ylim([0.0, 1.05]) + plt.xlabel('False Positive Rate') + plt.ylabel('True Positive Rate') + plt.title('Receiver Operating Characteristic') + plt.legend(loc="lower right") + plt.show() + + return roc_auc +# ---------------------------------------------------------------------------------------------------------- \ No newline at end of file diff --git a/ctlearn/core/pytorch/visualization/visualization_tsne.py b/ctlearn/core/pytorch/visualization/visualization_tsne.py new file mode 100644 index 00000000..51989406 --- /dev/null +++ b/ctlearn/core/pytorch/visualization/visualization_tsne.py @@ -0,0 +1,683 @@ +# from sklearn.manifold import TSNE +from cuml.manifold import TSNE + +# from nets.models.ResNet import ResNet +from ctlearn.ctlearn_helper.common import read_configuration +from ctlearn import utils +import torch +from tqdm import tqdm +import matplotlib.pyplot as plt +import numpy as np +from ctlearn.ctlearn_helper.CTlearnEnums import Task, Mode, EventType +from ctlearn.cli.run_model import Runner +from ctlearn.utils.utils import ( + load_pickle, + create_key_value_array + +) +from ctlearn import ( + CTADataset, + read_configuration, + ModelHelper, +) +# from threadpoolctl import threadpool_limits +from sklearn.decomposition import PCA +from sklearn.preprocessing import MinMaxScaler +from matplotlib import cm +import faiss +import pickle +from typing import Dict +import matplotlib +matplotlib.use('TkAgg') # Or 'Qt5Agg', depending on your setup + +def plt_tsne(embeddings, pred_labels, gt_labels, class_energy, classes, plot_3d=False, post_fix="tsne"): + """ + Visualize t-SNE embeddings with energy levels encoded by color, preserving class markers, + and highlighting misclassified points. + + Parameters: + - embeddings: t-SNE reduced feature vectors. + - pred_labels: Predicted class labels. + - gt_labels: Ground truth class labels. + - class_energy: Energy values corresponding to each sample. + - classes: List of class names. + - plot_3d: Boolean to plot in 3D. + - prefix: Prefix for the saved plot filename. + """ + + # Normalize energy values to [0, 1] for colormap mapping + scaler = MinMaxScaler() + norm_energy = scaler.fit_transform(class_energy.reshape(-1, 1)).flatten() + + # Choose a colormap + cmap = matplotlib.colormaps['viridis']#cm.get_cmap('viridis') + + # Identify misclassified points + misclassified = pred_labels != gt_labels + + # Initialize plot + fig = plt.figure(figsize=(10, 8)) + if plot_3d: + ax = fig.add_subplot(111, projection="3d") + else: + ax = fig.add_subplot(111) + + # Define markers for each class + markers = ["o", "^", "s"] + + # Plot correctly classified points with energy-based colors + for i, class_name in enumerate(classes): + idx = (gt_labels == i) & ~misclassified + if plot_3d: + sc = ax.scatter( + embeddings[idx, 0], + embeddings[idx, 1], + embeddings[idx, 2], + color=cmap(norm_energy[idx]), + marker=markers[i], + label=f"Correct {class_name}", + alpha=0.6 + ) + else: + sc = ax.scatter( + embeddings[idx, 0], + embeddings[idx, 1], + color=cmap(norm_energy[idx]), + marker=markers[i], + label=f"Correct {class_name}", + alpha=0.6 + ) + + # Plot misclassified points with a distinct marker and color + if plot_3d: + ax.scatter( + embeddings[misclassified, 0], + embeddings[misclassified, 1], + embeddings[misclassified, 2], + color="red", + marker="x", + label="Misclassified", + alpha=0.1 + ) + else: + ax.scatter( + embeddings[misclassified, 0], + embeddings[misclassified, 1], + color="red", + marker="x", + label="Misclassified", + alpha=0.1 + ) + + # Set plot title and labels + ax.legend(loc="best") + ax.set_title("t-SNE Visualization with Energy Levels and Misclassifications") + ax.set_xlabel("t-SNE 1") + ax.set_ylabel("t-SNE 2") + if plot_3d: + ax.set_zlabel("t-SNE 3") + + # Add colorbar to indicate energy levels + sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=class_energy.min(), vmax=class_energy.max())) + sm.set_array([]) + cbar = plt.colorbar(sm) + + cbar.set_label('Energy Levels') + + # Save and display plot + plt.savefig(f'./tsne_{post_fix}.png') + plt.show() + +def main(use_pickle,reduce_similarity,plot_tsne): + + classes_list = ["gamma", "proton"] + config_file_str = "./config/training_config_iaa_neutron_missing_analisys.yml" + parameters = read_configuration(config_file_str) + # file_name = parameters["data"]["train_gamma_proton"] + file_name = "/storage/ctlearn_data/training/node_pickles/proton_proton_theta_16.087_az_108.090_runs1-416_train_275107.dl1.pickle" + if not use_pickle: + + batch_size = 64 #parameters["hyp"]["batches"] + num_workers = parameters["dataset"]["num_workers"] + pin_memory = parameters["dataset"]["pin_memory"] + persistent_workers = parameters["dataset"]["persistent_workers"] + + print("Data loaded...") + + validation_data =load_pickle(file_name) + + factor = 1 + data_len = len(validation_data["data"]) + validation_data["data"] = validation_data["data"][0 : int(data_len / factor)] + validation_data["true_shower_primary_id"] = validation_data[ + "true_shower_primary_id" + ][0 : int(data_len / factor)] + + cta_ds_validation = CTADataset( + pickle_data=validation_data, task=Task.type,mode=Mode.results, parameters=parameters, use_augmentation = False + ) + + data_loader = torch.utils.data.DataLoader( + cta_ds_validation, + batch_size=batch_size, + shuffle=False, + collate_fn=cta_ds_validation.collate_fn, + num_workers=num_workers, + pin_memory=pin_memory, + persistent_workers=persistent_workers, + ) + + + runner = Runner() + model_net = runner.create_model(parameters["model"]["model_type"]) + + device_str = parameters["arch"]["device"] + device = torch.device(device_str) + + check_point_path = parameters["data"]["type_checkpoint"] + + model_net = ModelHelper.loadModel( + model_net, "", check_point_path, mode=Mode.results , device_str=device_str + ) + + model_net.eval() + + class_feature_vector = [] + class_predictions_class = [] + class_gt_class = [] + class_energy = [] + class_hillas = [] + class_hillas_name = validation_data["hillas_name"] + pbar = tqdm(total=len(data_loader), desc="DL2 conv", leave=True) + + + for batch_idx, (features, labels) in enumerate(data_loader): + + imgs = features["image"].to(device).contiguous() + peak_time = features["peak_time"].to(device).contiguous() + classification_pred, energy_pred, direction_pred = model_net( + imgs, peak_time + ) + + if len(features)==0: + continue + + classification_pred_ = classification_pred[0] + feature_vector = classification_pred[1].cpu().detach().numpy() + predicted = torch.softmax(classification_pred_, dim=1) + predicted_class = predicted.argmax(dim=1) + predicted_class = predicted_class.cpu().detach().numpy() + + labels_class = (labels["particletype"].int().to(device).contiguous()) + labels_class = labels_class.cpu().detach().numpy() + + labels_energy = labels["energy"].cpu().detach().numpy() + + hillas = features["hillas"] + hillas = { + key: tensor.cpu().detach().numpy() for key, tensor in hillas.items() + } + id = list(range(features["image"].shape[0])) + + hillas_vector = np.array(create_key_value_array(hillas, id)).T + + class_feature_vector.extend(feature_vector[:, :]) + class_predictions_class.extend(predicted_class[:]) + class_gt_class.extend(labels_class[:]) + energy = np.power(10,labels_energy[:])[:,0] + class_energy.extend(energy) + class_hillas.extend(hillas_vector) + if batch_idx % 10 == 0: + pbar.update(10) + + + embeddings = np.array(list(class_feature_vector)) + pred_labels = np.array(list(class_predictions_class)) + gt_labels = np.array(list(class_gt_class)) + gt_energies = np.array(list(class_energy)) + class_hillas = np.array(list(class_hillas)) + + data_dict ={"embeddings":embeddings, + "pred_labels":pred_labels, + "gt_labels":gt_labels, + "gt_energies":gt_energies, + "hillas_name":class_hillas_name, + "hillas":class_hillas} + + + # Save pickle + save_file_name= "./prediction_train.pickle" + with open(save_file_name, "wb") as handle: + pickle.dump(data_dict, handle, protocol=pickle.HIGHEST_PROTOCOL) + print("Files saved ... ", save_file_name) + + + save_file_name= "./prediction_train.pickle" + # save_file_name= "./validation_data.pickle" + + prediction = load_pickle(save_file_name) + print("pickle loaded...") + embeddings = prediction["embeddings"] + pred_labels = prediction["pred_labels"] + gt_labels = prediction["gt_labels"] + gt_energies = prediction["gt_energies"] + hillas_name = prediction["hillas_name"] + hillas = prediction["hillas"] + if reduce_similarity: + original_data =load_pickle(file_name) + SIMIL_THRS = 0.99 + + #------------------------------------------------------------------------------------------ + # dimension = embeddings.shape[1] + # nlist = 100 # Número de clusters (ajustar según dataset) + # # Crear el índice con clusters + # faiss.normalize_L2(embeddings) + # quantizer = faiss.IndexFlatIP(dimension) + # index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT) + # # index.nprobe = 20 + # index.train(embeddings) # Entrenar el índice de clustering + # index.add(embeddings) # Agregar embeddings + # if not index.is_trained: + # print("Index NOT trained properly.") + # exit() + # else: + # print("Index trained properly.") + #------------------------------------------------------------------------------------------ + # dimension = embeddings.shape[1] + # nlist = 100 # Número de clusters (ajustar según dataset) + # # Crear el índice con clusters + # faiss.normalize_L2(embeddings) + # # index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT) + # index = faiss.index_factory(dimension, f"IVF{nlist},Flat", faiss.METRIC_INNER_PRODUCT) + + # index.nprobe = 20 + # index.train(embeddings) # Entrenar el índice de clustering + # index.add(embeddings) # Agregar embeddings + # if not index.is_trained: + # print("Index NOT trained properly.") + # exit() + # else: + # print("Index trained properly.") + #------------------------------------------------------------------------------------------ + # dimension = embeddings.shape[1] + # faiss.normalize_L2(embeddings) + # index = faiss.IndexHNSWFlat(dimension, 32) # 32 vecinos en el grafo + # index.hnsw.efConstruction = 60 # Controla la calidad del grafo (más grande = mejor recall) + # index.hnsw.efSearch = 100 # Cuántos vecinos considerar en búsqueda + + # # Agregar embeddings + # index.add(embeddings) + #------------------------------------------------------------------------------------------ + # + #------------------------------------------------------------------------------------------ + # dimension = embeddings.shape[1] + # index = faiss.IndexFlatIP(dimension) # Index con producto interno (similaridad coseno) + # index = faiss.IndexIDMap(index) + # faiss.normalize_L2(embeddings) # Normalizar embeddings para similitud del coseno + # index.add(embeddings) # Agregar embeddings al índice de Faiss + #------------------------------------------------------------------------------------------ + # print("Generating index.") + # embeddings = np.array(embeddings).astype('float32') + # faiss.normalize_L2(embeddings) + # dimension = embeddings.shape[1] + + # nlist = 50 #100 # Número de clusters (ajustar según dataset) + # # Step 1: Create the CPU index + # quantizer = faiss.IndexFlatIP(dimension) + # cpu_index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT) + + # # Step 2: Train on CPU + # cpu_index.train(embeddings) + + # # Step 3: Move to GPU + # res = faiss.StandardGpuResources() # Use default options + # index = faiss.index_cpu_to_gpu(res, 0, cpu_index) # 0 is the GPU id + # index.nprobe = 80 + # # Step 4: Add embeddings to GPU index + # index.add(embeddings) + # if not index.is_trained: + # print("Index NOT trained properly.") + # exit() + # else: + # print("Index trained properly.") + + + print("Generating index.") + embeddings = np.array(embeddings).astype('float32') + faiss.normalize_L2(embeddings) + dimension = embeddings.shape[1] + + nlist = 50 # Número de clusters + quantizer = faiss.IndexFlatIP(dimension) + cpu_index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT) + + # Entrenar índice en CPU + cpu_index.train(embeddings) + num_gpus = faiss.get_num_gpus() + print("GPUs detected:",num_gpus) + # Crear recursos para múltiples GPUs + gpu_resources = [faiss.StandardGpuResources() for _ in range(num_gpus)] + + # Distribuir el índice entrenado a múltiples GPUs + gpu_indices = [ + faiss.index_cpu_to_gpu(gpu_resources[i], i, cpu_index) + for i in range(2) + ] + + # Combinar índices GPU en un índice shard + index = faiss.IndexShards(dimension, True, False) + for sub_index in gpu_indices: + index.add_shard(sub_index) + + index.nprobe = 80 + + # Agregar embeddings al índice distribuido (se distribuyen automáticamente entre GPUs) + index.add(embeddings) + + if not index.is_trained: + print("Index NOT trained properly.") + exit() + else: + print("Index trained properly.") + + #------------------------------------------------------------------------------------------ + + # print("Generating index.") + # embeddings = np.array(embeddings).astype('float32') + # faiss.normalize_L2(embeddings) + # dimension = embeddings.shape[1] + + # nlist = 50#100 # Número de clusters (ajustar según dataset) + + # # Paso 1: Crear el índice en CPU + # quantizer = faiss.IndexFlatIP(dimension) + # index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT) + + # # Paso 2: Entrenar el índice (CPU) + # index.train(embeddings) + + # index.nprobe = 80 + + # # Paso 3: Añadir los embeddings al índice CPU + # index.add(embeddings) + + # # Verificación del entrenamiento + # if not index.is_trained: + # print("Index NOT trained properly.") + # exit() + # else: + # print("Index trained properly.") + #------------------------------------------------------------------------------------------ + + print("Looking for duplicates: Search...") + # D, I = index.search(embeddings, k=3) # Search for the k nearest + # lims, D, I = index.range_search(embeddings, SIMIL_THRS) + # Parameters + SIMIL_THRS = 0.99 + k_search_default= 10 + k_search_first_bins = 45 #25 + similarity_default = 0.95 + similarity_first_bins = 0.90 + k_bin=11 + cut_off_leakage_intensity = 0.2 + cut_off_intensity = 50 + bins = np.logspace(np.log10(2.51e-02), 2, 19) + bin_indices = np.digitize(gt_energies, bins) + num_bins = len(bins) # should be 18 in your case + + # Create array with default value 0.99 + simil_thresholds = np.full(num_bins+1, similarity_default) + k_search = np.full(num_bins+1, k_search_default,dtype=int) + k_search[:k_bin] = k_search_first_bins + # Set first 5 bins to 0.95 + simil_thresholds[:k_bin] = similarity_first_bins + # === Deduplication within bins === + keep_image = set() + deleted = set() + deleted_bin = set() + seen = set() + + unique_bins = np.unique(bin_indices) + print(f"Processing {len(unique_bins)} bins...") + + deleted_list=[] + for b in tqdm(unique_bins): + # Get indices in the current bin + bin_mask = (bin_indices == b) + bin_ids = np.where(bin_mask)[0] + + if len(bin_ids) < 2: + continue # Skip bins with fewer than 2 items + + # Extract embeddings for this bin + bin_embeddings = embeddings[bin_ids] + print(f"Processing bin {b} of {len(unique_bins)} ") + # Perform FAISS k-NN search + D, I = index.search(bin_embeddings, k=int(k_search[b])) + # D, I = index.search(bin_embeddings, k=k_search_default) + + + for i in range(len(bin_ids)): + query_idx = bin_ids[i] + + if query_idx in deleted: + continue + + if hillas[query_idx][hillas_name["leakage_pixels_width_2"]]>cut_off_leakage_intensity or hillas[query_idx][hillas_name["hillas_intensity"]]= len(bin_ids): + continue # Index out of bounds (can happen if few items in bin) + + neighbor_idx = bin_ids[neighbor_idx_local] + + if neighbor_idx == query_idx: + continue + + if similarity >= simil_thresholds[b] and neighbor_idx not in seen: + deleted.add(neighbor_idx) + deleted_bin.add(neighbor_idx) + seen.add(neighbor_idx) + + deleted_list.append(deleted_bin) + deleted_bin = set() + seen = set() + print(f"Final samples kept: {len(keep_image)}") + print(f"Samples removed as similar: {len(deleted)}") + + # Count number of deletions per bin + deleted_counts = [len(bin_deleted) for bin_deleted in deleted_list] + kept_indices = sorted(keep_image) + + #------------------------------------------------------------------------------------------ + # Save data + # original_data["data"]=original_data["data"][kept_indices] + original_data["data"] = [original_data["data"][i] for i in kept_indices] + save_file_name= "./train_reduced_data.pickle" + + with open(save_file_name, "wb") as handle: + pickle.dump(original_data, handle, protocol=pickle.HIGHEST_PROTOCOL) + #------------------------------------------------------------------------------------------ + # Get indices for each class in the ground truth + gamma_indices = np.where(gt_labels == 0)[0] + proton_indices = np.where(gt_labels == 1)[0] + + # Gamma accuracy: correct predictions among all gamma ground truths + gamma_accuracy = np.mean(pred_labels[gamma_indices] == 0) + + # Proton accuracy: correct predictions among all proton ground truths + proton_accuracy = np.mean(pred_labels[proton_indices] == 1) + + overall_accuracy = np.mean(gt_labels == pred_labels) + + print(f"Accuracy Before: {overall_accuracy}") + print(f"Accuracy Gamma Before: {gamma_accuracy}") + print(f"Accuracy Proton Before: {proton_accuracy}") + print(f"Num of Gammas Before: {len(gamma_indices)}") + print(f"Num of Protons Before: {len(proton_indices)}") + #------------------------------------------------------------------------------------------ + # Get indices for each class in the ground truth + gamma_indices = np.where(gt_labels[kept_indices] == 0)[0] + proton_indices = np.where(gt_labels[kept_indices] == 1)[0] + + # Gamma accuracy: correct predictions among all gamma ground truths + gamma_accuracy = np.mean(pred_labels[kept_indices][gamma_indices] == 0) + + # Proton accuracy: correct predictions among all proton ground truths + proton_accuracy = np.mean(pred_labels[kept_indices][proton_indices] == 1) + + overall_accuracy = np.mean(gt_labels[kept_indices] == pred_labels[kept_indices]) + + print(f"Accuracy After: {overall_accuracy}") + print(f"Accuracy Gamma After: {gamma_accuracy}") + print(f"Accuracy Proton After: {proton_accuracy}") + print(f"Num of Gammas After: {len(gamma_indices)}") + print(f"Num of Protons After: {len(proton_indices)}") + #------------------------------------------------------------------------------------------ + # Plot + plt.figure(figsize=(10, 5)) + plt.bar(range(len(deleted_counts)), deleted_counts) + plt.xlabel("Bin Index") + plt.ylabel("Number of Deleted Samples") + plt.title("Deleted Samples per Energy Bin") + plt.grid(True) + plt.tight_layout() + plt.savefig(f'./histogram_deleted_{k_search_first_bins}_{k_search_default}.png') + # plt.show() + + #------------------------------------------------------------------------------------------ + # Create bin labels (log-scale bins) + # Ensure we have 18 bins + # Bin labels: 18 labels for 18 bins + # num_bins = len(bins) + # bin_labels = [f"{bins[i]:.2e}-{bins[i+1]:.2e}" for i in range(num_bins-1)] + # deleted_counts = [len(s) for s in deleted_list] + # x = list(range(num_bins)) + + # plt.figure(figsize=(12, 6)) + # plt.bar(x, deleted_counts) + # plt.xticks(ticks=x, labels=bin_labels, rotation=45, ha='right') + # plt.xlabel("Energy Bin") + # plt.ylabel("Number of Deleted Samples") + # plt.title("Deleted Samples per Energy Bin") + # plt.grid(True, which='both', linestyle='--', linewidth=0.5) + # plt.tight_layout() + # plt.show() + #------------------------------------------------------------------------------------------ + # Convert sets to sorted lists (for consistent indexing) + deleted_indices = sorted(deleted) + kept_indices = sorted(keep_image) + + + # Get energy values + energies_before = gt_energies + energies_after = gt_energies[kept_indices] + + # Create log-spaced bins for histogram (same as your binning) + hist_bins = np.logspace(np.log10(2.51e-02), 2, 30) + + plt.figure(figsize=(12, 6)) + + # Plot before as outline + plt.hist(energies_before, bins=hist_bins, histtype='step', label='Before Filtering', color='gray', linewidth=1.5) + + # Plot after as filled + plt.hist(energies_after, bins=hist_bins, alpha=0.6, label='After Filtering', color='green') + + plt.xscale('log') + plt.xlabel('Energy') + plt.ylabel('Number of Samples') + plt.title('Energy Histogram Before and After Filtering') + plt.legend() + plt.grid(True, which='both', ls='--', linewidth=0.5) + plt.tight_layout() + plt.savefig(f'./histogram_energy_before_after_{k_search_first_bins}_{k_search_default}.png') + # plt.show() + #------------------------------------------------------------------------------------------ + + # flat_similarities = D[:, 1:].flatten() # excluye el self-match + # plt.hist(flat_similarities, bins=100) + # plt.axvline(SIMIL_THRS, color='red', linestyle='--') + # plt.title("Distribución de Similitudes (excluyendo self-match)") + # plt.xlabel("Similitud") + # plt.ylabel("Frecuencia") + # plt.show() + + + # EPS = 1e-5 + # keep_image = set() + # deleted = set() + # seen = set() + + # for i in tqdm(range(len(embeddings))): + # if i in deleted: + # continue + + # keep_image.add(i) + # if gt_energies[i]<1: + # for j in range(1, len(I[i])): + # neighbor_idx = I[i][j] + # similarity = D[i][j] + + # if similarity < (SIMIL_THRS - EPS): + # continue # no lo consideres duplicado + + # if neighbor_idx not in keep_image: + # deleted.add(neighbor_idx) + #------------------------------------------------------------------------------------------ + # keep_image = set() + # deleted = set() + # print("Search Done") + # for i in tqdm(range(len(embeddings))): + # if i in deleted: + # continue # Si ya está marcada como duplicada, ignorarla + # keep_image.add(i) # Mantener esta imagen + # start_idx = lims[i] + # end_idx = lims[i + 1] + # for j in range(start_idx, end_idx): + # neighbor_idx = I[j] + # if neighbor_idx != i: # Evitar la propia imagen + # deleted.add(neighbor_idx) # Marcar como duplicada + # # print(f"Eliminando duplicado: {image_paths[neighbor_idx]} (ID {image_ids[neighbor_idx]})") + + + print(f"keeped: {len(keep_image)}") + print(f"deleted: {len(deleted)}") + embeddings = embeddings[list(keep_image)] + pred_labels = pred_labels[list(keep_image)] + gt_labels = gt_labels[list(keep_image)] + gt_energies = gt_energies[list(keep_image)] + + if plot_tsne: + # energy_threshold = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] + # energy_threshold = [2, 3, 4, 5, 6, 7, 8, 9, 10] + print("Plotting t-sne") + + energy_threshold=[0] + for energy_thr in energy_threshold: + embeddings_ = embeddings[gt_energies>energy_thr] + pred_labels_ = pred_labels[gt_energies>energy_thr] + gt_labels_ = gt_labels[gt_energies>energy_thr] + gt_energies_ = gt_energies[gt_energies>energy_thr] + + + print(f"generating thr:{energy_threshold}") + plt_tsne(embeddings_, pred_labels_, gt_labels_, gt_energies_, classes_list, post_fix="v4_"+str(energy_thr)) + + +if __name__ == "__main__": + + use_pickle = True + reduce_similarity = True + plot_tsne=True + + main(use_pickle,reduce_similarity,plot_tsne) diff --git a/ctlearn/core/tests/test_attention.py b/ctlearn/core/tests/test_attention.py new file mode 100644 index 00000000..072a9279 --- /dev/null +++ b/ctlearn/core/tests/test_attention.py @@ -0,0 +1,62 @@ +import pytest +pytest.importorskip("keras") +import numpy as np +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_pytorch.py similarity index 53% rename from ctlearn/core/tests/test_loader.py rename to ctlearn/core/tests/test_loader_pytorch.py index 7fe71900..47fbe8cf 100644 --- a/ctlearn/core/tests/test_loader.py +++ b/ctlearn/core/tests/test_loader_pytorch.py @@ -1,7 +1,13 @@ +import pytest +pytest.importorskip("dl1_data_handler") +torch = pytest.importorskip("torch") from traitlets.config.loader import Config from dl1_data_handler.reader import DLImageReader -from ctlearn.core.loader import DLDataLoader +from ctlearn.core.data_loader.loader import DLDataLoader + +from ctlearn.tools.train.pytorch.utils import read_configuration +from ctlearn.tools.train.pytorch.utils import get_absolute_config_path def test_data_loader(dl1_gamma_file): @@ -18,14 +24,28 @@ def test_data_loader(dl1_gamma_file): # Create an image reader dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) # Create a data loader - dl1_loader = DLDataLoader( + config_file_dir = get_absolute_config_path() + print(config_file_dir) + parameters = read_configuration(config_file_dir) + + dl1_loader = DLDataLoader.create( + framework = "pytorch", DLDataReader=dl1_reader, indices=[0], tasks=["type", "energy", "cameradirection", "skydirection"], batch_size=1, + parameters=parameters, + use_augmentation=parameters["augmentation"]["use_augmentation"], + is_training=True ) # Get the features and labels fgrom the data loader for one batch - features, labels = dl1_loader[0] + print(len(dl1_loader[0])) + print(dl1_loader[0][0]) + print(dl1_loader[0][1]) + print(dl1_loader[0][2]) + + + features, labels, _ = dl1_loader[0] # Check that all the correct labels are present assert ( "type" in labels @@ -34,4 +54,7 @@ def test_data_loader(dl1_gamma_file): and "skydirection" in labels ) # Check the shape of the features - assert features.shape == (1, 110, 110, 2) + assert features["image"].shape == (1, 2, 110, 110) + +if __name__ == "__main__": + test_data_loader() \ No newline at end of file diff --git a/ctlearn/core/tests/test_models.py b/ctlearn/core/tests/test_models.py new file mode 100644 index 00000000..6eaa1962 --- /dev/null +++ b/ctlearn/core/tests/test_models.py @@ -0,0 +1,574 @@ +import pytest +pytest.importorskip("keras") +import re +import keras +import numpy as np +import pytest +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}" + ) + + +''' +TODO: Do this correctly with more time +def _copy_head_weights(keras_wrapper, torch_wrapper, tasks): + """Shared helper to copy Dense -> Linear head weights for mapped tasks.""" + for task in tasks: + k_dense_layers = [ + l 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] + p_linear_modules = [m for m in p_head_module if isinstance(m, torch.nn.Linear)] + + for k_layer, p_module in zip(k_dense_layers, p_linear_modules): + weights = k_layer.get_weights() + if not weights: + continue + w = np.transpose(weights[0], (1, 0)) # Keras (In, Out) -> PyTorch (Out, In) + p_module.weight.data = torch.from_numpy(w).float() + if len(weights) > 1 and p_module.bias is not None: + p_module.bias.data = torch.from_numpy(weights[1]).float() + +def _copy_weights_single_cnn(keras_wrapper, torch_wrapper, tasks): + """Transfers weights between KerasSingleCNN and PyTorchSingleCNN sequentially.""" + k_backbone = keras_wrapper.backbone_model + p_backbone = torch_wrapper.backbone_model + + # 1. Include Dense / Linear layers alongside Conv and BN + def extract_keras_layers(model): + found = [] + for l in model.layers: + # Recurse into nested sub-models / custom layers if present + if hasattr(l, "layers") and not isinstance(l, (keras.layers.Conv2D, keras.layers.Dense, keras.layers.BatchNormalization)): + found.extend(extract_keras_layers(l)) + elif isinstance(l, (keras.layers.Conv2D, keras.layers.Dense, keras.layers.BatchNormalization)) and len(l.weights) > 0: + found.append(l) + return found + + k_layers = extract_keras_layers(k_backbone) + + # 2. Match corresponding PyTorch parameter-bearing modules + p_modules = [ + m for m in p_backbone.modules() + if isinstance(m, (torch.nn.Conv2d, torch.nn.Linear, torch.nn.BatchNorm2d)) + ] + + # Guard against architecture/sequence mismatch + if len(k_layers) != len(p_modules): + raise ValueError( + f"Layer count mismatch! Keras found {len(k_layers)} parameter layers, " + f"PyTorch found {len(p_modules)}. Check if Linear/Dense or SE modules are missing." + ) + + for k_layer, p_module in zip(k_layers, p_modules): + weights = k_layer.get_weights() + if not weights: + continue + + # --- Handle Conv2D / Conv2d --- + if isinstance(k_layer, keras.layers.Conv2D): + assert isinstance(p_module, torch.nn.Conv2d), f"Type mismatch: {type(k_layer)} vs {type(p_module)}" + # Keras Conv2D: (H, W, In, Out) -> PyTorch Conv2d: (Out, In, H, W) + w = np.transpose(weights[0], (3, 2, 0, 1)) + + assert p_module.weight.shape == w.shape, f"Shape mismatch in Conv2D: Keras transposed {w.shape} vs PyTorch {p_module.weight.shape}" + p_module.weight.data = torch.from_numpy(w).float() + + if len(weights) > 1 and p_module.bias is not None: + p_module.bias.data = torch.from_numpy(weights[1]).float() + elif p_module.bias is not None: + p_module.bias.data.zero_() + + # --- Handle Dense / Linear (Crucial for Channel-SE & Dual-SE!) --- + elif isinstance(k_layer, keras.layers.Dense): + assert isinstance(p_module, torch.nn.Linear), f"Type mismatch: {type(k_layer)} vs {type(p_module)}" + # Keras Dense: (In, Out) -> PyTorch Linear: (Out, In) + w = np.transpose(weights[0], (1, 0)) + + assert p_module.weight.shape == w.shape, f"Shape mismatch in Dense: Keras transposed {w.shape} vs PyTorch {p_module.weight.shape}" + p_module.weight.data = torch.from_numpy(w).float() + + if len(weights) > 1 and p_module.bias is not None: + p_module.bias.data = torch.from_numpy(weights[1]).float() + elif p_module.bias is not None: + p_module.bias.data.zero_() + + # --- Handle Batch Normalization --- + elif isinstance(k_layer, keras.layers.BatchNormalization): + assert isinstance(p_module, torch.nn.BatchNorm2d), f"Type mismatch: {type(k_layer)} vs {type(p_module)}" + # Keras BN weights order: [gamma (scale), beta (bias), moving_mean, moving_variance] + p_module.weight.data = torch.from_numpy(weights[0]).float() + p_module.bias.data = torch.from_numpy(weights[1]).float() + p_module.running_mean.data = torch.from_numpy(weights[2]).float() + p_module.running_var.data = torch.from_numpy(weights[3]).float() + + _copy_head_weights(keras_wrapper, torch_wrapper, tasks) + +def _copy_weights_resnet(keras_wrapper, torch_wrapper, tasks): + """Transfers weights between KerasResNet and PyTorchResNet block by block, + + including BatchNorm, Dense (SE), and Conv layers. + """ + k_backbone = keras_wrapper.backbone_model + p_backbone = torch_wrapper.backbone_model + + # 1. Transfer Initial Layers (Conv + BN) + k_init_layers = [ + l for l in k_backbone.layers + if isinstance(l, (keras.layers.Conv2D, keras.layers.BatchNormalization)) + and not re.search(r"conv\d+_block\d+", l.name) + ] + p_init_modules = [ + m for m in p_backbone.children() + if isinstance(m, (torch.nn.Conv2d, torch.nn.BatchNorm2d)) + ] + + for k_l, p_m in zip(k_init_layers, p_init_modules): + _transfer_single_layer_weights(k_l, p_m) + + # 2. Collect PyTorch Residual Blocks + p_res_blocks = [ + m for m in p_backbone.modules() + if isinstance(m, (BasicBlock, BottleneckBlock)) + ] + + # 3. Group Keras layers by block key + block_pattern = re.compile(r"(conv\d+_block\d+)") + keras_blocks_dict = {} + for layer in k_backbone.layers: + match = block_pattern.search(layer.name) + if match: + block_key = match.group(1) + keras_blocks_dict.setdefault(block_key, []).append(layer) + + def _sort_key(key_str): + numbers = re.findall(r"\d+", key_str) + return int(numbers[0]), int(numbers[1]) + + sorted_keras_keys = sorted(keras_blocks_dict.keys(), key=_sort_key) + + # 4. Sequential Intra-Block Transfer (Handles Conv, BN, and Dense / SE layers) + for block_key, p_block in zip(sorted_keras_keys, p_res_blocks): + k_block_layers = [ + l for l in keras_blocks_dict[block_key] + if len(l.weights) > 0 + ] + p_block_modules = [ + m for m in p_block.modules() + if isinstance(m, (torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.Linear)) + ] + + if len(k_block_layers) != len(p_block_modules): + raise ValueError( + f"Block {block_key} mismatch! Keras has {len(k_block_layers)} weighted layers, " + f"PyTorch block has {len(p_block_modules)} modules." + ) + + for k_layer, p_module in zip(k_block_layers, p_block_modules): + _transfer_single_layer_weights(k_layer, p_module) + + # 5. Transfer Task Head Weights + _copy_head_weights(keras_wrapper, torch_wrapper, tasks) + + +def _transfer_single_layer_weights(k_layer, p_module): + """Helper to copy parameter weights from a Keras layer to a PyTorch module.""" + weights = k_layer.get_weights() + if not weights: + return + + if isinstance(k_layer, keras.layers.Conv2D): + assert isinstance(p_module, torch.nn.Conv2d) + w = np.transpose(weights[0], (3, 2, 0, 1)) + p_module.weight.data = torch.from_numpy(w).float() + if len(weights) > 1 and p_module.bias is not None: + p_module.bias.data = torch.from_numpy(weights[1]).float() + + elif isinstance(k_layer, keras.layers.Dense): + if isinstance(p_module, torch.nn.Linear): + # Dense -> Linear + w = np.transpose(weights[0], (1, 0)) # (in_features, out_features) -> (out_features, in_features) + p_module.weight.data = torch.from_numpy(w).float() + if len(weights) > 1 and p_module.bias is not None: + p_module.bias.data = torch.from_numpy(weights[1]).float() + + elif isinstance(p_module, torch.nn.Conv2d): + # Dense -> 1x1 Conv2d (SE Block representation) + # Keras Dense weight shape: (in_features, out_features) + # PyTorch Conv2d weight shape: (out_channels, in_channels, 1, 1) + w = np.transpose(weights[0], (1, 0))[:, :, None, None] + p_module.weight.data = torch.from_numpy(w).float() + if len(weights) > 1 and p_module.bias is not None: + p_module.bias.data = torch.from_numpy(weights[1]).float() + else: + raise TypeError(f"Unsupported PyTorch module type {type(p_module)} for Keras Dense layer.") + + elif isinstance(k_layer, keras.layers.BatchNormalization): + assert isinstance(p_module, torch.nn.BatchNorm2d) + p_module.weight.data = torch.from_numpy(weights[0]).float() # gamma + p_module.bias.data = torch.from_numpy(weights[1]).float() # beta + p_module.running_mean.data = torch.from_numpy(weights[2]).float() # mean + p_module.running_var.data = torch.from_numpy(weights[3]).float() # variance + + +@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_numerical_outputs_with_aligned_weights(common_config, batchnorm, attention): + """Verify that models yield similar predictions once weights are copied.""" + 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 + ) + _copy_weights_single_cnn(keras_wrapper, torch_wrapper, [task]) + + # Prepare identical input data + x_keras = rng.standard_normal(size=(2, *common_config["input_shape_keras"])).astype(np.float32) + x_torch = torch.from_numpy(np.transpose(x_keras, (0, 3, 1, 2))) + # Evaluate models in eval mode + keras_preds = keras_wrapper.model(x_keras, training=False) + torch_wrapper.model.eval() + with torch.no_grad(): + torch_preds = torch_wrapper.model(x_torch) + + # Extract predicted value array for the task + if isinstance(keras_preds, dict): + k_tensor = keras_preds[task] + else: + k_tensor = keras_preds + + k_val = k_tensor.numpy() if hasattr(k_tensor, "numpy") else np.array(k_tensor) + + # Handle PyTorch prediction output + if isinstance(torch_preds, dict): + p_val = torch_preds[task].cpu().numpy() + else: + p_val = torch_preds.cpu().numpy() + + # Compare outputs within tight numerical tolerance + np.testing.assert_allclose( + k_val, + p_val, + atol=0.025, + err_msg=f"Value divergence detected in task '{task}'", + ) + + +@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"] + absolute_tolerance = { + "basic": 0.1, + "bottleneck": 0.05, + } + + 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 + ) + + # Copy weights for ResNet + _copy_weights_resnet(keras_wrapper, torch_wrapper, [task]) + + x_keras = rng.standard_normal(size=(2, *common_config["input_shape_keras"])).astype(np.float32) + x_torch = torch.from_numpy(np.transpose(x_keras, (0, 3, 1, 2))) + + keras_preds = keras_wrapper.model(x_keras, training=False) + torch_wrapper.model.eval() + with torch.no_grad(): + torch_preds = torch_wrapper.model(x_torch) + + k_val = keras_preds[task].numpy() if hasattr(keras_preds[task], "numpy") else np.array(keras_preds[task]) + p_val = torch_preds[task].cpu().numpy() + + np.testing.assert_allclose( + k_val, + p_val, + atol=absolute_tolerance[block_type], + err_msg=f"Value divergence detected in ResNet ({block_type}) for task '{task}'", + ) +''' \ No newline at end of file diff --git a/ctlearn/tools/__init__.py b/ctlearn/tools/__init__.py index d54df32d..c6d237fc 100644 --- a/ctlearn/tools/__init__.py +++ b/ctlearn/tools/__init__.py @@ -1,13 +1,37 @@ """ctlearn command line tools. """ -from .train_model import TrainCTLearnModel -from .predict_LST1 import LST1PredictionTool -from .predict_model import MonoPredictCTLearnModel, StereoPredictCTLearnModel +import sys +import os +import warnings + +is_debug = '--debug' in sys.argv or any(arg.startswith('--log-level=DEBUG') for arg in sys.argv) +if not is_debug: + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' + os.environ['NCCL_DEBUG'] = 'WARN' + warnings.filterwarnings("ignore", ".*NoneDefaultNotAllowedWarning.*") + warnings.filterwarnings("ignore", ".*MergeConflictWarning.*") + warnings.filterwarnings("ignore", ".*'ctlearn.tools.train_model' found in sys.modules.*") + +from .train_model import DLFrameWork +try: + from .predict_LST1 import LST1PredictionTool +except ImportError: + pass +try: + from .predict_model import MonoPredictCTLearnModel, StereoPredictCTLearnModel +except ImportError: + pass __all__ = [ - "TrainCTLearnModel", - "LST1PredictionTool", - "MonoPredictCTLearnModel", - "StereoPredictCTLearnModel" -] \ No newline at end of file + "DLFrameWork", +] +try: + __all__.append("MonoPredictCTLearnModel") + __all__.append("StereoPredictCTLearnModel") +except NameError: + pass +try: + __all__.append("LST1PredictionTool") +except NameError: + pass \ No newline at end of file diff --git a/ctlearn/tools/_version.py b/ctlearn/tools/_version.py new file mode 100644 index 00000000..b8023d8b --- /dev/null +++ b/ctlearn/tools/_version.py @@ -0,0 +1 @@ +__version__ = '0.0.1' diff --git a/ctlearn/tools/conftest.py b/ctlearn/tools/conftest.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/tools/predict/keras/predic_LST1_keras.py b/ctlearn/tools/predict/keras/predic_LST1_keras.py new file mode 100644 index 00000000..7118a67f --- /dev/null +++ b/ctlearn/tools/predict/keras/predic_LST1_keras.py @@ -0,0 +1,255 @@ +""" +Keras prediction module for LST1 telescope data. +This module provides functionality to load trained Keras models and perform predictions +on DL1 level data for particle type classification, energy estimation, and direction reconstruction. +""" + +from ctapipe.io import read_table +from astropy.table import join +import keras +from dl1_data_handler.reader import get_unmapped_image +import numpy as np + + +def predictions(self): + """ + Perform predictions on input DL1 data using trained Keras models. + + This function processes the input file in batches, applies quality cuts, + and generates predictions for particle type, energy, and/or direction + depending on the configured models. The models are split into backbone + and head components to extract feature vectors. + + Returns + ------- + tuple + Contains the following arrays: + - event_id: Event identifiers + - tel_azimuth: Telescope azimuth angles + - tel_altitude: Telescope altitude angles + - trigger_time: Event trigger times in MJD + - prediction: Particle type classification scores (gammaness) + - energy: Reconstructed energy values + - cam_coord_offset_x: Camera coordinate offset in x direction + - cam_coord_offset_y: Camera coordinate offset in y direction + - classification_fvs: Classification feature vectors from backbone + - energy_fvs: Energy estimation feature vectors from backbone + - direction_fvs: Direction reconstruction feature vectors from backbone + + Notes + ----- + The function processes data in batches to manage memory efficiently and + applies quality selection criteria before making predictions. + """ + # Initialize output arrays for storing results + event_id, tel_azimuth, tel_altitude, trigger_time = [], [], [], [] + prediction, energy, cam_coord_offset_x, cam_coord_offset_y = [], [], [], [] + classification_fvs, energy_fvs, direction_fvs = [], [], [] + + # Process input file in batches + for start in range(0, self.table_length, self.batch_size): + stop = min(start + self.batch_size, self.table_length) + self.log.debug("Processing chunk from '%d' to '%d'.", start, stop - 1) + + # Read the DL1 data table for current batch + dl1_table = read_table( + self.input_url, self.image_table_path, start=start, stop=stop + ) + + # Join tables to enable quality selection + # Join with parameter table for event parameters + dl1_table = join( + left=dl1_table, + right=self.parameter_table, + keys=["event_id"], + ) + # Join with trigger table for timing information + dl1_table = join( + left=dl1_table, + right=self.trigger_table, + keys=["event_id"], + ) + + # Apply quality selection criteria + # Initialize mask to accept all events initially + passes_quality_checks = np.ones(len(dl1_table), dtype=bool) + + # Apply quality query if configured + if self.quality_query: + passes_quality_checks = self.quality_query.get_table_mask(dl1_table) + + # Filter events based on quality criteria + dl1_table = dl1_table[passes_quality_checks] + + # Skip batch if no events passed quality selection + if len(dl1_table) == 0: + self.log.debug("No events passed the quality selection.") + continue + + # Prepare input data by mapping images to model input format + data = [] + for event in dl1_table: + # Get the unmapped image with specified channels and transforms + image = get_unmapped_image(event, self.channels, self.transforms) + # Map image to model's expected input format + data.append(self.image_mapper.map_image(image)) + input_data = {"input": np.array(data)} + + # Handle compatibility between Keras 2 and Keras 3 + # Keras 3 expects direct array input, not dictionary + if int(keras.__version__.split(".")[0]) >= 3: + input_data = input_data["input"] + + # Store event metadata + event_id.extend(dl1_table["event_id"].data) + tel_azimuth.extend(dl1_table["tel_az"].data) + tel_altitude.extend(dl1_table["tel_alt"].data) + trigger_time.extend(dl1_table["time"].mjd) + + # Perform particle type classification if model is loaded + if self.load_type_model_from is not None: + # Extract feature vectors from backbone + classification_feature_vectors = self.backbone_type.predict_on_batch(input_data) + classification_fvs.extend(classification_feature_vectors) + # Generate predictions from head model + predict_data = self.head_type.predict_on_batch(classification_feature_vectors) + # Extract gammaness score (probability of being gamma) + prediction.extend(predict_data[:, 1]) + + # Perform energy estimation if model is loaded + if self.load_energy_model_from is not None: + # Extract feature vectors from backbone + energy_feature_vectors = self.backbone_energy.predict_on_batch(input_data) + energy_fvs.extend(energy_feature_vectors) + # Generate energy predictions from head model + predict_data = self.head_energy.predict_on_batch(energy_feature_vectors) + energy.extend(predict_data.T[0]) + + # Perform direction reconstruction if model is loaded + if self.load_cameradirection_model_from is not None: + # Extract feature vectors from backbone + direction_feature_vectors = self.backbone_direction.predict_on_batch(input_data) + direction_fvs.extend(direction_feature_vectors) + # Generate direction predictions from head model + predict_data = self.head_direction.predict_on_batch(direction_feature_vectors) + # Extract x and y components of camera coordinate offset + cam_coord_offset_x.extend(predict_data.T[0]) + cam_coord_offset_y.extend(predict_data.T[1]) + + return (event_id, tel_azimuth, tel_altitude, trigger_time, prediction, energy, + cam_coord_offset_x, cam_coord_offset_y, classification_fvs, energy_fvs, direction_fvs) + + +def _split_model(model): + """ + Split a Keras model into backbone and head components. + + This function separates a trained model into two parts: + - Backbone: Feature extraction layers (typically convolutional layers) + - Head: Task-specific prediction layers (typically dense layers) + + This separation allows extraction of intermediate feature representations + which can be useful for analysis or transfer learning. + + Parameters + ---------- + model : keras.Model + Complete trained Keras model to be split. The model should have: + - Layer 0: Input layer + - Layer 1: Backbone (feature extractor) + - Layers 2+: Head (prediction layers) + + Returns + ------- + backbone : keras.Model + Feature extraction model that outputs intermediate representations. + head : keras.Model + Prediction model that takes backbone outputs and produces final predictions. + + Notes + ----- + The function assumes a specific model architecture where the backbone + is the second layer (index 1) of the complete model. This is a common + pattern in CTLearn models where the backbone is wrapped as a single layer. + """ + # Extract the backbone model (second layer of the complete model) + # Layer 0 is the input, layer 1 is the backbone feature extractor + backbone = model.get_layer(index=1) + + # Create a new head model using layers after the backbone + # Define input with the same shape as backbone output + backbone_output_shape = keras.Input(model.layers[2].input_shape[1:]) + x = backbone_output_shape + + # Reconstruct head by connecting all layers after backbone + for layer in model.layers[2:]: + x = layer(x) + + # Create the head model + head = keras.Model(inputs=backbone_output_shape, outputs=x) + + return backbone, head + + +def load_keras_model(self): + """ + Load Keras models from saved files and split them into backbone and head. + + This function loads trained Keras models for different tasks (particle type + classification, energy estimation, direction reconstruction) and splits each + into backbone and head components for efficient prediction and feature extraction. + + Parameters + ---------- + self : object + Prediction handler instance containing model paths: + - load_type_model_from: Path to particle classification model + - load_energy_model_from: Path to energy estimation model + - load_cameradirection_model_from: Path to direction reconstruction model + + Returns + ------- + input_shape : tuple + Shape of the model input (height, width, channels). + Returns the shape from the last loaded model. + + Notes + ----- + The function sets the following attributes on self: + - backbone_type, head_type: Split models for particle classification + - backbone_energy, head_energy: Split models for energy estimation + - backbone_direction, head_direction: Split models for direction reconstruction + """ + input_shape = None + + # Load particle type classification model if configured + 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:] + # Split model into backbone and head + self.backbone_type, self.head_type = _split_model(model_type) + + # Load energy estimation model if configured + 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:] + # Split model into backbone and head + self.backbone_energy, self.head_energy = _split_model(model_energy) + + # Load direction reconstruction model if configured + 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:] + # Split model into backbone and head + self.backbone_direction, self.head_direction = _split_model(model_direction) + + return input_shape diff --git a/ctlearn/tools/predict/keras/predic_model_keras.py b/ctlearn/tools/predict/keras/predic_model_keras.py new file mode 100644 index 00000000..0d2b3472 --- /dev/null +++ b/ctlearn/tools/predict/keras/predic_model_keras.py @@ -0,0 +1,156 @@ +""" +Keras model prediction module for CTLearn. +This module provides functionality to load trained Keras models and perform batch predictions +on DL1 data, with optional feature vector extraction from backbone models. +""" + +from ctlearn.core.data_loader.loader import DLDataLoader +import keras +from astropy.table import Table, vstack +import numpy as np + + +def predict_with_model(self, model_path, task): + """ + Load and predict with a CTLearn Keras model. + + This function loads a trained model from the specified path and performs predictions + on the provided data. It handles both complete and incomplete batches, and optionally + extracts feature vectors from the backbone model for downstream analysis. + + Parameters + ---------- + model_path : str + Path to a Keras model file (Keras 3) or directory (Keras 2). + The model should be a complete trained CTLearn model. + task : str + The task for which prediction is being made (e.g., 'type', 'energy', 'cameradirection'). + + Returns + ------- + predict_data : astropy.table.Table + Table containing the prediction results with columns corresponding to + the model's output. + feature_vectors : np.ndarray or None + Feature vectors extracted from the backbone model if dl1_features is enabled. + Returns None if feature extraction is not requested. + """ + # Create data loader for the main batch processing + # The DLDataLoader is initialized separately for each prediction task + # Batch size is multiplied by number of replicas for distributed inference + data_loader = DLDataLoader.create( + framework="keras", + DLDataReader=self.dl1dh_reader, + indices=self.indices, + tasks=[], + batch_size=self.batch_size * self.strategy.num_replicas_in_sync, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + + # Handle incomplete last batch + # Keras only processes complete batches during prediction, so we need + # a separate data loader for remaining events that don't fill a complete batch + data_loader_last_batch = None + if self.last_batch_size > 0: + # Extract indices for the last incomplete batch + last_batch_indices = self.indices[-self.last_batch_size:] + data_loader_last_batch = DLDataLoader.create( + framework="keras", + DLDataReader=self.dl1dh_reader, + indices=last_batch_indices, + tasks=[], + batch_size=self.last_batch_size, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + + # Load the trained model from the specified path + model = keras.saving.load_model(model_path) + + # Initialize variables for optional feature extraction + backbone_model, feature_vectors = None, None + + if self.dl1_features: + # Feature extraction mode: split model into backbone and head + # This allows us to extract intermediate representations (feature vectors) + + # Extract the backbone model (second layer of the complete model) + # Layer 0: Input, Layer 1: Backbone (feature extractor), Layers 2+: Head + backbone_model = model.get_layer(index=1) + + # Reconstruct the head model from layers after the backbone + # Define input with the same shape as backbone output + backbone_output_shape = keras.Input(model.layers[2].input_shape[1:]) + x = backbone_output_shape + + # Connect all layers after backbone to create head model + for layer in model.layers[2:]: + x = layer(x) + head = keras.Model(inputs=backbone_output_shape, outputs=x) + + # Extract feature vectors from backbone + feature_vectors = backbone_model.predict( + data_loader, verbose=self.keras_verbose + ) + + # Generate predictions from head using extracted features + predict_data = Table( + { + task: head.predict( + feature_vectors, verbose=self.keras_verbose + ) + } + ) + + # Process last incomplete batch if it exists + if data_loader_last_batch is not None: + # Extract features from last batch + feature_vectors_last_batch = backbone_model.predict( + data_loader_last_batch, verbose=self.keras_verbose + ) + # Concatenate feature vectors from all batches + feature_vectors = np.concatenate( + (feature_vectors, feature_vectors_last_batch) + ) + # Generate predictions for last batch and stack with main predictions + predict_data = vstack( + [ + predict_data, + Table( + { + task: head.predict( + feature_vectors_last_batch, + verbose=self.keras_verbose, + ) + } + ), + ] + ) + else: + # Standard prediction mode without feature extraction + # Use the complete model for end-to-end prediction + predict_data = model.predict(data_loader, verbose=self.keras_verbose) + + # Convert predictions to Astropy Table + if isinstance(predict_data, dict): + predict_data = Table(predict_data) + else: + predict_data = Table({task: predict_data}) + + # Process last incomplete batch if it exists + if data_loader_last_batch is not None: + # Generate predictions for last batch + predict_data_last_batch = model.predict( + data_loader_last_batch, verbose=self.keras_verbose + ) + + if isinstance(predict_data_last_batch, dict): + predict_data_last_batch = Table(predict_data_last_batch) + else: + predict_data_last_batch = Table({task: predict_data_last_batch}) + + # Stack predictions from main batches and last batch + predict_data = vstack([predict_data, predict_data_last_batch]) + + return predict_data, feature_vectors \ No newline at end of file diff --git a/ctlearn/tools/predict/predict_LST1.py b/ctlearn/tools/predict/predict_LST1.py new file mode 100644 index 00000000..ab3671bc --- /dev/null +++ b/ctlearn/tools/predict/predict_LST1.py @@ -0,0 +1,834 @@ +""" +Predict the gammaness, energy and arrival direction from lstchain DL1 data. +""" + +import numpy as np +import tables +import keras +from astropy import units as u +from astropy.coordinates import AltAz,SkyCoord +from astropy.table import Table, setdiff, vstack +from astropy.coordinates import AltAz,SkyCoord +from astropy.table import Table, setdiff, vstack +from astropy.time import Time + +from ctapipe.containers import ( + ParticleClassificationContainer, + ReconstructedGeometryContainer, + ReconstructedEnergyContainer, +) +from ctapipe.coordinates import CameraFrame +from ctapipe.core import Tool +from ctapipe.core.tool import ToolConfigurationError +from ctapipe.core.traits import ( + Bool, + Int, + Path, + List, + CaselessStrEnum, + ComponentName, + Unicode, + UseEnum, + classes_with_traits, +) +from ctapipe.instrument.optics import FocalLengthKind +from ctapipe.io import read_table, write_table +from ctapipe.reco.utils import add_defaults_and_meta + +from ctlearn.utils import get_lst1_subarray_description +from dl1_data_handler.image_mapper import ImageMapper +from dl1_data_handler.reader import TableQualityQuery +from ctlearn.tools.predict.utils.load_model import load_model +from ctlearn.core.ctlearn_enum import Task, Mode +from ctlearn.tools.train.pytorch.utils import ( + sanity_check, + read_configuration, + expected_structure, +) + +POINTING_GROUP = "/dl1/monitoring/telescope/pointing" +DL1_TELESCOPE_GROUP = "/dl1/event/telescope" +DL2_TELESCOPE_GROUP = "/dl2/event/telescope" +DL2_SUBARRAY_GROUP = "/dl2/event/subarray" +SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] +TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] + + +class LST1PredictionTool(Tool): + """ + Tool to predict the gammaness, energy and arrival direction from lstchain DL1 data. + + This tool is used to predict the gammaness, energy and arrival direction + from pixel-wise image data in lstchain format. The tool loads the trained models + from the specified paths and performs inference on the input data. The + input data is expected to be in the DL1 format of lstchain and the output data is + stored in the DL2 format of ctapipe. Besides the DL2 predictions, the tool creates + the SubarrayDescription of the LST-1 telescope and stores it in the output file. + In addition, the tool also creates the trigger, pointing and DL1 parameters tables + and stores them in the output file. + + CAUTION: The tool is designed to work with the DL1 data format of lstchain only. + """ + + name = "LST1PredictionTool" + description = __doc__ + + examples = """ + To predict from DL1 lstchain data using trained CTLearn models: + > ctlearn-predict-model \\ + --input_url input.subrun.lstchain.dl1.h5 \\ + --LST1PredictionTool.batch_size=64 \\ + --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" \\ + --output output.dl2.h5 \\ + --overwrite \\ + """ + + input_url = Path( + help="Input LST-1 HDF5 files including pixel-wise image data", + allow_none=True, + exists=True, + directory_ok=False, + file_ok=True, + ).tag(config=True) + + prefix = Unicode( + default_value="CTLearn", + allow_none=False, + help="Name of the reconstruction algorithm used to generate the dl2 data.", + ).tag(config=True) + + load_type_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) " + "for the classification of the primary particle type." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_energy_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) " + "for the regression of the primary particle energy." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_cameradirection_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) " + "for the regression of the primary particle arrival direction " + "based on the camera coordinate offsets." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + batch_size = Int( + default_value=128, + allow_none=False, + help="Size of the batch to perform inference of the neural network.", + ).tag(config=True) + + channels = List( + trait=CaselessStrEnum( + [ + "image", + "cleaned_image", + "peak_time", + "relative_peak_time", + "cleaned_peak_time", + "cleaned_relative_peak_time", + ] + ), + default_value=["cleaned_image", "cleaned_peak_time"], + allow_none=False, + help=( + "Set the input channels to be loaded from the DL1 event data. " + "image: integrated charges, " + "cleaned_image: integrated charges cleaned with the DL1 cleaning mask, " + "peak_time: extracted peak arrival times, " + "relative_peak_time: extracted relative peak arrival times, " + "cleaned_peak_time: extracted peak arrival times cleaned with the DL1 cleaning mask, " + "cleaned_relative_peak_time: extracted relative peak arrival times cleaned with the DL1 cleaning mask." + ), + ).tag(config=True) + + image_mapper_type = ComponentName(ImageMapper, default_value="BilinearMapper").tag( + config=True + ) + + focal_length_choice = UseEnum( + FocalLengthKind, + default_value=FocalLengthKind.EFFECTIVE, + help=( + "If both nominal and effective focal lengths are available, " + " which one to use for the `~ctapipe.coordinates.CameraFrame` attached" + " to the `~ctapipe.instrument.CameraGeometry` instances in the" + " `~ctapipe.instrument.SubarrayDescription` which will be used in" + " CameraFrame to TelescopeFrame coordinate transforms." + " The 'nominal' focal length is the one used during " + " the simulation, the 'effective' focal length is computed using specialized " + " ray-tracing from a point light source" + ), + ).tag(config=True) + + override_obs_id = Int( + default_value=None, + allow_none=True, + help=( + "Use the given obs_id instead of the default one. " + "Needed to merge subruns later with ctapipe-merge." + ), + ).tag(config=True) + + output_path = Path( + default_value="./output.dl2.h5", + allow_none=False, + help="Output path to save the dl2 prediction results", + ).tag(config=True) + + pytorch_config_file = Path( + default_value="./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml", + help="Pytorch config file", + ).tag(config=True) + + framework_type = CaselessStrEnum( + ["pytorch", "keras"], + default_value="keras", + help="Framework to use: pytorch or keras", + ).tag(config=True) + + overwrite = Bool(help="Overwrite output file if it exists").tag(config=True) + + aliases = { + ("i", "input_url"): "LST1PredictionTool.input_url", + ("t", "type_model"): "LST1PredictionTool.load_type_model_from", + ("e", "energy_model"): "LST1PredictionTool.load_energy_model_from", + ("d", "cameradirection_model"): "LST1PredictionTool.load_cameradirection_model_from", + ("o", "output"): "LST1PredictionTool.output_path", + ("f", "framework"): "LST1PredictionTool.framework_type", + ("p", "pytorch_config_file"): "LST1PredictionTool.pytorch_config_file", + + } + + flags = { + "overwrite": ( + {"LST1PredictionTool": {"overwrite": True}}, + "Overwrite existing files", + ), + } + + classes = classes_with_traits(ImageMapper) + + def _predictions(self): + if self.framework_type == "keras": + self.log.info("Using the Keras Model") + from ctlearn.tools.predict.keras.predic_LST1_keras import predictions + return predictions(self) + + elif self.framework_type == "pytorch": + self.log.info("Using the Pytorch Model") + from ctlearn.tools.predict.pytorch.predic_LST1_pytorch import predictions + return predictions(self) + + def setup(self): + # Save dl1 image and parameters tree schemas and tel id for easy access + import torch + self.image_table_path = "/dl1/event/telescope/image/LST_LSTCam" + self.parameter_table_name = "/dl1/event/telescope/parameters/LST_LSTCam" + self.tel_id = 1 + if self.framework_type == "pytorch": + self.log.info(f"Using {self.pytorch_config_file} config file for pytorch framework") + self.parameters = read_configuration(self.pytorch_config_file) + sanity_check(self.parameters, expected_structure) + self.batch_size = self.parameters["hyp"]["batches"] + self.device_str = self.parameters["arch"]["device"] + self.optim_batch_size = self.parameters["hyp"]["dynamic_batches"] + self.device = torch.device(self.device_str) + self.tasks = [] + self.type_mu = self.parameters["normalization"]["type_mu"] + self.type_sigma = self.parameters["normalization"]["type_sigma"] + self.dir_mu = self.parameters["normalization"]["dir_mu"] + self.dir_sigma = self.parameters["normalization"]["dir_sigma"] + self.energy_mu = self.parameters["normalization"]["energy_mu"] + self.energy_sigma = self.parameters["normalization"]["energy_sigma"] + + if self.load_type_model_from is not None: + self.tasks.append(Task.type) + if self.load_energy_model_from is not None: + self.tasks.append(Task.energy) + if self.load_cameradirection_model_from is not None: + self.tasks.append(Task.direction) + + # 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 + input_shape = load_model(self) + + # 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 + self.subarray.to_hdf(self.output_path, overwrite=self.overwrite) + self.log.info("SubarrayDescription was stored in '%s'", self.output_path) + # Initialize the Table data quality query + self.quality_query = TableQualityQuery(parent=self) + # Copy the pixel rotation of the camera geometry of the subarray in a variable + # since the ImageMapper will be derotated the pixels. The pixel rotation + # is needed to create a rotated camera frame in order to transform the + # predicted camera coordinate offsets back to the correct Alt/Az coordinates. + self.pix_rotation = self.subarray.tel[self.tel_id].camera.geometry.pix_rotation + # Create the ImageMapper + self.image_mapper = ImageMapper.from_name( + name=self.image_mapper_type, + geometry=self.subarray.tel[self.tel_id].camera.geometry, + subarray=self.subarray, + parent=self, + ) + # Check if the input shape of the model matches the image shape of the ImageMapper + if self.framework_type == "keras": + 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]}' ." + ) + + # Get offset and scaling of images + self.transforms = {} + self.transforms["image_scale"] = 0.0 + self.transforms["image_offset"] = 0 + self.transforms["peak_time_scale"] = 0.0 + self.transforms["peak_time_offset"] = 0 + + # 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[ + "CTAFIELD_3_TRANSFORM_SCALE" + ] + self.transforms["image_offset"] = img_table_v_attrs[ + "CTAFIELD_3_TRANSFORM_OFFSET" + ] + if "CTAFIELD_4_TRANSFORM_SCALE" in img_table_v_attrs: + self.transforms["peak_time_scale"] = img_table_v_attrs[ + "CTAFIELD_4_TRANSFORM_SCALE" + ] + self.transforms["peak_time_offset"] = img_table_v_attrs[ + "CTAFIELD_4_TRANSFORM_OFFSET" + ] + + 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: + all_identifiers["obs_id"] = self.override_obs_id + self.obs_id = all_identifiers["obs_id"][0] + self.parameter_table = all_identifiers.copy() + tel_az = u.Quantity(self.parameter_table["az_tel"], unit=u.rad) + tel_alt = u.Quantity(self.parameter_table["alt_tel"], unit=u.rad) + event_type = self.parameter_table["event_type"] + time = Time(self.parameter_table["dragon_time"] * u.s, format="unix") + # Create the pointing table + # This table is used to store the telescope pointing per event + pointing_table = Table( + { + "time": time, + "azimuth": tel_az, + "altitude": tel_alt, + } + ) + write_table( + pointing_table, + self.output_path, + f"{POINTING_GROUP}/tel_{self.tel_id:03d}", + overwrite=self.overwrite, + ) + self.log.info( + "DL1 telescope pointing table was stored in '%s' under '%s'", + self.output_path, + f"{POINTING_GROUP}/tel_{self.tel_id:03d}", + ) + # Set the time format to MJD since in the other table we store the time in MJD + time.format = "mjd" + # Keep only the necessary columns for the creation of tables + all_identifiers.keep_columns(TELESCOPE_EVENT_KEYS) + + # Create the dl1 telescope trigger table + self.trigger_table = all_identifiers.copy() + self.trigger_table.add_column(time, name="time") + self.trigger_table.add_column(-1, name="n_trigger_pixels") + + write_table( + self.trigger_table, + self.output_path, + "/dl1/event/telescope/trigger", + overwrite=self.overwrite, + ) + self.log.info( + "DL1 telescope trigger table was stored in '%s' under '%s'", + self.output_path, + "/dl1/event/telescope/trigger", + ) + self.trigger_table.keep_columns(["obs_id", "event_id", "time"]) + self.trigger_table.add_column( + np.ones((len(self.trigger_table), 1), dtype=bool), name="tel_with_trigger" + ) + self.trigger_table.add_column(event_type, name="event_type") + + # Save the dl1 subrray trigger table to the output file + # write_table( + # self.trigger_table, + # self.output_path, + # "/dl1/event/subarray/trigger", + # overwrite=self.overwrite, + # ) + # self.log.info( + # "DL1 subarray trigger table was stored in '%s' under '%s'", + # self.output_path, + # "/dl1/event/subarray/trigger", + # ) + # Create the dl1 parameters table + self.parameter_table.rename_column("intensity", "hillas_intensity") + self.parameter_table.rename_column("x", "hillas_x") + self.parameter_table.rename_column("y", "hillas_y") + self.parameter_table.rename_column("phi", "hillas_phi") + self.parameter_table.rename_column("psi", "hillas_psi") + self.parameter_table.rename_column("length", "hillas_length") + self.parameter_table.rename_column("length_uncertainty", "hillas_length_uncertainty") + self.parameter_table.rename_column("width", "hillas_width") + self.parameter_table.rename_column("width_uncertainty", "hillas_width_uncertainty") + self.parameter_table.rename_column("skewness", "hillas_skewness") + self.parameter_table.rename_column("kurtosis", "hillas_kurtosis") + self.parameter_table.rename_column("time_gradient", "timing_deviation") + self.parameter_table.rename_column("intercept", "timing_intercept") + self.parameter_table.rename_column("n_pixels", "morphology_n_pixels") + self.parameter_table.rename_column("n_islands", "morphology_n_islands") + self.parameter_table.keep_columns( + [ + "obs_id", + "event_id", + "hillas_intensity", + "hillas_x", + "hillas_y", + "hillas_phi", + "hillas_psi", + "hillas_length", + "hillas_length_uncertainty", + "hillas_width", + "hillas_width_uncertainty", + "hillas_skewness", + "hillas_kurtosis", + "timing_deviation", + "timing_intercept", + "morphology_n_pixels", + "morphology_n_islands", + ] + ) + self.parameter_table.add_column(self.tel_id, name="tel_id", index=2) + # Save the dl1 parameters table to the output file + write_table( + self.parameter_table, + self.output_path, + f"/dl1/event/telescope/parameters/tel_{self.tel_id:03d}", + overwrite=self.overwrite, + ) + self.log.info( + "DL1 parameters table was stored in '%s' under '%s'", + self.output_path, + f"/dl1/event/telescope/parameters/tel_{self.tel_id:03d}", + ) + + # Add additional columns to the parameter table + # which are not present in the originl DL1 parameter table. + # They are needed for applying the quality selection. + self.parameter_table.add_column(event_type, name="event_type") + self.parameter_table.add_column(tel_az, name="tel_az") + self.parameter_table.add_column(tel_alt, name="tel_alt") + # Only select cosmic events for the prediction + self.parameter_table = self.parameter_table[self.parameter_table["event_type"]==32] + + self.log.info("Starting the prediction...") + # Iterate over the data in chunks based on the batch size + event_id, tel_azimuth, tel_altitude, trigger_time, prediction, energy, cam_coord_offset_x, cam_coord_offset_y, classification_fvs, energy_fvs, direction_fvs = self._predictions() + + # Create the prediction tables + example_identifiers = Table( + { + "obs_id": np.full(len(event_id), self.obs_id, dtype=int), + "event_id": event_id, + "tel_id": np.full(len(event_id), self.tel_id, dtype=int), + } + ) + nonexample_identifiers = setdiff( + all_identifiers, example_identifiers, keys=TELESCOPE_EVENT_KEYS + ) + if len(nonexample_identifiers) > 0: + nonexample_identifiers.sort(TELESCOPE_EVENT_KEYS) + # Create the feature vector table + feature_vector_table = example_identifiers.copy() + fvs_columns_list, fvs_shapes_list = [], [] + if self.load_type_model_from is not None: + classification_table = example_identifiers.copy() + classification_table.add_column( + prediction, name=f"{self.prefix}_tel_prediction" + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = self._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_prediction"], + shapes=[(len(nonexample_identifiers),)], + ) + classification_table = vstack([classification_table, nan_table]) + classification_table.sort(TELESCOPE_EVENT_KEYS) + classification_is_valid = ~np.isnan(classification_table[f"{self.prefix}_tel_prediction"].data, dtype=bool) + classification_table.add_column( + classification_is_valid, + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + classification_table, + ParticleClassificationContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + # Save the prediction to the output file + write_table( + classification_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{self.tel_id:03d}", + overwrite=self.overwrite, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{self.tel_id:03d}", + ) + # Write the mono telescope prediction to the subarray prediction table + subarray_classification_table = classification_table.copy() + subarray_classification_table.remove_column("tel_id") + for colname in subarray_classification_table.colnames: + if "_tel_" in colname: + subarray_classification_table.rename_column( + colname, colname.replace("_tel", "") + ) + subarray_classification_table.add_column( + classification_is_valid, name=f"{self.prefix}_telescopes" + ) + # Save the prediction to the output file + write_table( + subarray_classification_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + overwrite=self.overwrite, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + ) + # Adding the feature vectors for the classification + is_valid_col = ~np.isnan( + np.min(classification_fvs, axis=1), dtype=bool + ) + feature_vector_table.add_column( + classification_fvs, + name=f"{self.prefix}_tel_classification_feature_vectors", + ) + if nonexample_identifiers is not None: + fvs_columns_list.append(f"{self.prefix}_tel_classification_feature_vectors") + fvs_shapes_list.append( + ( + len(nonexample_identifiers), + classification_fvs[0].shape[0], + ) + ) + if self.load_energy_model_from is not None: + energy_table = example_identifiers.copy() + # Convert the reconstructed energy from log10(TeV) to TeV + reco_energy = u.Quantity(np.power(10, np.squeeze(energy)), unit=u.TeV) + # Add the reconstructed energy to the prediction table + energy_table.add_column(reco_energy, name=f"{self.prefix}_tel_energy") + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = self._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_energy"], + shapes=[(len(nonexample_identifiers),)], + ) + energy_table = vstack([energy_table, nan_table]) + energy_table.sort(TELESCOPE_EVENT_KEYS) + energy_is_valid = ~np.isnan(energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool) + energy_table.add_column( + energy_is_valid, + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + energy_table, + ReconstructedEnergyContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + # Save the prediction to the output file + write_table( + energy_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{self.tel_id:03d}", + overwrite=self.overwrite, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{self.tel_id:03d}", + ) + # Write the mono telescope prediction to the subarray prediction table + subarray_energy_table = energy_table.copy() + subarray_energy_table.remove_column("tel_id") + for colname in subarray_energy_table.colnames: + if "_tel_" in colname: + subarray_energy_table.rename_column( + colname, colname.replace("_tel", "") + ) + subarray_energy_table.add_column( + energy_is_valid, name=f"{self.prefix}_telescopes" + ) + # Save the prediction to the output file + write_table( + subarray_energy_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + overwrite=self.overwrite, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + ) + # Adding the feature vectors for the energy regression + is_valid_col = ~np.isnan( + np.min(energy_fvs, axis=1), dtype=bool + ) + feature_vector_table.add_column( + energy_fvs, + name=f"{self.prefix}_tel_energy_feature_vectors", + ) + if nonexample_identifiers is not None: + fvs_columns_list.append(f"{self.prefix}_tel_energy_feature_vectors") + fvs_shapes_list.append( + ( + len(nonexample_identifiers), + energy_fvs[0].shape[0], + ) + ) + if self.load_cameradirection_model_from is not None: + direction_table = example_identifiers.copy() + # Set the telescope position + tel_ground_frame = self.subarray.tel_coords[ + self.subarray.tel_ids_to_indices(self.tel_id) + ] + # Set the telescope pointing with the trigger timestamp and the telescope position + trigger_time = Time(trigger_time, format="mjd") + altaz = AltAz( + location=tel_ground_frame.to_earth_location(), + obstime=trigger_time, + ) + # Set the telescope pointing + tel_pointing = SkyCoord( + az=u.Quantity(tel_azimuth, unit=u.rad), + alt=u.Quantity(tel_altitude, unit=u.rad), + frame=altaz, + ) + # Set a new camera frame with the pixel rotation of the camera + camera_frame = CameraFrame( + focal_length=self.subarray.tel[self.tel_id].camera.geometry.frame.focal_length, + rotation=self.pix_rotation, + telescope_pointing=tel_pointing, + ) + ## Remove (save the cam_coord_offset_x, cam_coord_offset_y) predicted by the model in a pickel file + + # with open('/lhome/ext/ucm147/ucm1477/ctlearn/test_local_cristian/cam_coord_offset_test.pkl', 'wb') as f: + # import pickle + # pickle.dump((cam_coord_offset_x, cam_coord_offset_y), f) + # print('Pickell camera_coordinates saved') + # Set the camera coordinate offset + cam_coord_offset = SkyCoord( + x=u.Quantity(cam_coord_offset_x, unit=u.m), + y=u.Quantity(cam_coord_offset_y, unit=u.m), + frame=camera_frame + ) + # Transform the true Alt/Az coordinates to camera coordinates + reco_direction = cam_coord_offset.transform_to(altaz) + # Add the reconstructed direction (az, alt) to the prediction table + direction_table.add_column( + reco_direction.az.to(u.deg), name=f"{self.prefix}_tel_az" + ) + direction_table.add_column( + reco_direction.alt.to(u.deg), name=f"{self.prefix}_tel_alt" + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = self._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_az", f"{self.prefix}_tel_alt"], + shapes=[(len(nonexample_identifiers),), (len(nonexample_identifiers),)], + ) + direction_table = vstack([direction_table, nan_table]) + direction_table.keep_columns( + TELESCOPE_EVENT_KEYS + + [f"{self.prefix}_tel_az", f"{self.prefix}_tel_alt"] + ) + direction_table.sort(TELESCOPE_EVENT_KEYS) + direction_is_valid = ~np.isnan(direction_table[f"{self.prefix}_tel_az"].data, dtype=bool) + direction_table.add_column( + direction_is_valid, + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + direction_table, + ReconstructedGeometryContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + # Save the prediction to the output file + write_table( + direction_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{self.tel_id:03d}", + overwrite=self.overwrite, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{self.tel_id:03d}", + ) + # Write the mono telescope prediction to the subarray prediction table + subarray_direction_table = direction_table.copy() + subarray_direction_table.remove_column("tel_id") + for colname in subarray_direction_table.colnames: + if "_tel_" in colname: + subarray_direction_table.rename_column( + colname, colname.replace("_tel", "") + ) + subarray_direction_table.add_column( + direction_is_valid, name=f"{self.prefix}_telescopes" + ) + # Save the prediction to the output file + write_table( + subarray_direction_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + overwrite=self.overwrite, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + ) + # Adding the feature vectors for the arrival direction regression + is_valid_col = ~np.isnan( + np.min(direction_fvs, axis=1), dtype=bool + ) + feature_vector_table.add_column( + direction_fvs, + name=f"{self.prefix}_tel_direction_feature_vectors", + ) + if nonexample_identifiers is not None: + fvs_columns_list.append(f"{self.prefix}_tel_direction_feature_vectors") + fvs_shapes_list.append( + ( + len(nonexample_identifiers), + direction_fvs[0].shape[0], + ) + ) + # Produce output table with NaNs for missing predictions + if nonexample_identifiers is not None: + if len(nonexample_identifiers) > 0: + nan_table = self._create_nan_table( + nonexample_identifiers, + columns=fvs_columns_list, + shapes=fvs_shapes_list, + ) + feature_vector_table = vstack([feature_vector_table, nan_table]) + is_valid_col = np.concatenate( + (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) + ) + # Add is_valid column to the feature vector table + feature_vector_table.add_column( + is_valid_col, + name=f"{self.prefix}_tel_is_valid", + ) + # Save the prediction to the output file + write_table( + feature_vector_table, + self.output_path, + f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{self.tel_id:03d}", + overwrite=self.overwrite, + ) + self.log.info( + "DL1 feature vectors was stored in '%s' under '%s'", + self.output_path, + f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{self.tel_id:03d}", + ) + + def finish(self): + self.log.info("Tool is shutting down") + + def _create_nan_table(self, nonexample_identifiers, columns, shapes): + """ + Create a table with NaNs for missing predictions. + + This method creates a table with NaNs for missing predictions for the non-example identifiers. + + Parameters: + ----------- + nonexample_identifiers : astropy.table.Table + Table containing the non-example identifiers. + columns : list of str + List of column names to create in the table. + shapes : list of shapes + List of shapes for the columns to create in the table. + + Returns: + -------- + nan_table : astropy.table.Table + Table containing NaNs for missing predictions. + """ + # Create a table with NaNs for missing predictions + nan_table = nonexample_identifiers.copy() + for column_name, shape in zip(columns, shapes): + nan_table.add_column(np.full(shape, np.nan), name=column_name) + return nan_table + + +def main(): + # Run the tool + tool = LST1PredictionTool() + tool.run() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/ctlearn/tools/predict/predict_mono.py b/ctlearn/tools/predict/predict_mono.py new file mode 100644 index 00000000..86782cbd --- /dev/null +++ b/ctlearn/tools/predict/predict_mono.py @@ -0,0 +1,541 @@ +""" +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``. +""" + + +import numpy as np +from astropy.table import ( + Table, + vstack, + join, + setdiff, +) +from ctapipe.containers import ( + ParticleClassificationContainer, + ReconstructedGeometryContainer, + ReconstructedEnergyContainer, +) + +from ctapipe.core.traits import ComponentName +from ctapipe.io import write_table +from ctapipe.reco.reconstructor import ReconstructionProperty +from ctapipe.reco.stereo_combination import StereoCombiner +from ctapipe.reco.utils import add_defaults_and_meta +from dl1_data_handler.reader import ProcessType + + +SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" +FIXED_POINTING_GROUP = "/configuration/telescope/pointing" +POINTING_GROUP = "/dl1/monitoring/telescope/pointing" +SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" +DL1_TELESCOPE_GROUP = "/dl1/event/telescope" +DL1_SUBARRAY_GROUP = "/dl1/event/subarray" +DL2_SUBARRAY_GROUP = "/dl2/event/subarray" +DL2_TELESCOPE_GROUP = "/dl2/event/telescope" +SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] +TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] + +__all__ = ["MonoPredictCTLearnModel"] + +from ctlearn.tools.predict.utils.predict_model import PredictCTLearnModel + +class MonoPredictCTLearnModel(PredictCTLearnModel): + """ + Tool to predict the gammaness, energy and arrival direction from monoscopic R1/DL1 data using CTLearn models. + + This tool extends the ``PredictCTLearnModel`` to specifically handle monoscopic R1/DL1 data. The prediction + is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. + It also stores the telescope pointing monitoring and DL1 feature vectors (if selected) in the output file. + + Attributes + ---------- + name : str + Name of the tool. + description : str + Description of the tool. + examples : str + Examples of how to use the tool. + + Methods + ------- + start() + Start the tool. + _store_mc_telescope_pointing(all_identifiers) + Store the telescope pointing table for the mono mode for MC simulation. + """ + + name = "ctlearn-predict-mono-model" + description = __doc__ + + examples = """ + To predict from pixel-wise image data in mono mode using trained CTLearn models: + > ctlearn-predict-mono-model \\ + --input_url input.dl1.h5 \\ + --PredictCTLearnModel.batch_size=64 \\ + --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ + --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" \\ + --dl1-features \\ + --use-HDF5Merger \\ + --no-dl1-images \\ + --no-true-images \\ + --output output.dl2.h5 \\ + --PredictCTLearnModel.overwrite_tables=True \\ + + To predict from pixel-wise waveform data in mono mode using trained CTLearn models: + > ctlearn-predict-mono-model \\ + --input_url input.r1.h5 \\ + --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" \\ + --use-HDF5Merger \\ + --no-r0-waveforms \\ + --no-r1-waveforms \\ + --no-dl1-images \\ + --no-true-images \\ + --output output.dl2.h5 \\ + --PredictCTLearnModel.overwrite_tables=True \\ + """ + + stereo_combiner_cls = ComponentName( + StereoCombiner, + default_value="StereoMeanCombiner", + help="Which stereo combination method to use after the monoscopic reconstruction.", + ).tag(config=True) + + def start(self): + self.log.info("Processing the telescope pointings...") + # Retrieve the IDs from the dl1dh for the prediction tables + example_identifiers = self.dl1dh_reader.example_identifiers.copy() + example_identifiers.keep_columns(TELESCOPE_EVENT_KEYS) + all_identifiers = self.dl1dh_reader.tel_trigger_table.copy() + all_identifiers.keep_columns(TELESCOPE_EVENT_KEYS + ["time"]) + nonexample_identifiers = setdiff( + all_identifiers, example_identifiers, keys=TELESCOPE_EVENT_KEYS + ) + nonexample_identifiers.remove_column("time") + # Pointing table for the mono mode for MC simulation + if self.dl1dh_reader.process_type == ProcessType.Simulation: + pointing_info = self._store_mc_telescope_pointing(all_identifiers) + + # Pointing table for the observation mode + if self.dl1dh_reader.process_type == ProcessType.Observation: + pointing_info = super()._store_pointing(all_identifiers) + + self.log.info("Starting the prediction...") + classification_feature_vectors = None + if self.load_type_model_from is not None: + self.type_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefix, + property=ReconstructionProperty.PARTICLE_TYPE, + parent=self, + ) + # Predict the energy of the primary particle + classification_table, classification_feature_vectors = ( + super()._predict_classification(example_identifiers) + ) + if self.dl2_telescope: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_prediction"], + shapes=[(len(nonexample_identifiers),)], + ) + classification_table = vstack([classification_table, nan_table]) + # Add is_valid column to the energy table + classification_table.add_column( + ~np.isnan( + classification_table[f"{self.prefix}_tel_prediction"].data, + dtype=bool, + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + classification_table, + ParticleClassificationContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = classification_table["tel_id"] == tel_id + classification_tel_table = classification_table[telescope_mask] + classification_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Save the prediction to the output file for the selected telescope + write_table( + classification_tel_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info("Processing and storing the subarray type prediction...") + # Combine the telescope predictions to the subarray prediction using the stereo combiner + subarray_classification_table = self.type_stereo_combiner.predict_table( + classification_table + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if ( + subarray_classification_table[f"{self.prefix}_telescopes"].dtype + != np.bool_ + ): + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + ( + len(subarray_classification_table), + len(self.dl1dh_reader.tel_ids), + ), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + subarray_classification_table[f"{self.prefix}_telescopes"] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] = True + # Overwrite the column with the boolean mask with fix length + subarray_classification_table[f"{self.prefix}_telescopes"] = ( + reco_telescopes + ) + # Save the prediction to the output file + write_table( + subarray_classification_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + ) + energy_feature_vectors = None + if self.load_energy_model_from is not None: + self.energy_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefix, + property=ReconstructionProperty.ENERGY, + parent=self, + ) + # Predict the energy of the primary particle + energy_table, energy_feature_vectors = super()._predict_energy( + example_identifiers + ) + if self.dl2_telescope: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_energy"], + shapes=[(len(nonexample_identifiers),)], + ) + energy_table = vstack([energy_table, nan_table]) + # Add is_valid column to the energy table + energy_table.add_column( + ~np.isnan( + energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + energy_table, + ReconstructedEnergyContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = energy_table["tel_id"] == tel_id + energy_tel_table = energy_table[telescope_mask] + energy_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Save the prediction to the output file + write_table( + energy_tel_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info( + "Processing and storing the subarray energy prediction..." + ) + # Combine the telescope predictions to the subarray prediction using the stereo combiner + subarray_energy_table = self.energy_stereo_combiner.predict_table( + energy_table + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if subarray_energy_table[f"{self.prefix}_telescopes"].dtype != np.bool_: + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + (len(subarray_energy_table), len(self.dl1dh_reader.tel_ids)), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + subarray_energy_table[f"{self.prefix}_telescopes"] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] = True + # Overwrite the column with the boolean mask with fix length + subarray_energy_table[f"{self.prefix}_telescopes"] = reco_telescopes + # Save the prediction to the output file + write_table( + subarray_energy_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + ) + direction_feature_vectors = None + if self.load_cameradirection_model_from is not None: + self.geometry_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefix, + property=ReconstructionProperty.GEOMETRY, + parent=self, + ) + # Join the prediction table with the telescope pointing table + example_identifiers = join( + left=example_identifiers, + right=pointing_info, + keys=TELESCOPE_EVENT_KEYS, + ) + # Predict the arrival direction of the primary particle + direction_table, direction_feature_vectors = ( + super()._predict_cameradirection(example_identifiers) + ) + direction_tel_tables = [] + if self.dl2_telescope: + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = direction_table["tel_id"] == tel_id + direction_tel_table = direction_table[telescope_mask] + direction_tel_table = super()._transform_cam_coord_offsets_to_sky( + direction_tel_table + ) + # Produce output table with NaNs for missing predictions + nan_telescope_mask = nonexample_identifiers["tel_id"] == tel_id + nonexample_identifiers_tel = nonexample_identifiers[ + nan_telescope_mask + ] + if len(nonexample_identifiers_tel) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers_tel, + columns=[f"{self.prefix}_tel_alt", f"{self.prefix}_tel_az"], + shapes=[ + (len(nonexample_identifiers_tel),), + (len(nonexample_identifiers_tel),), + ], + ) + direction_tel_table = vstack([direction_tel_table, nan_table]) + direction_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Add is_valid column to the direction table + direction_tel_table.add_column( + ~np.isnan( + direction_tel_table[f"{self.prefix}_tel_alt"].data, + dtype=bool, + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + direction_tel_table, + ReconstructedGeometryContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + direction_tel_tables.append(direction_tel_table) + # Save the prediction to the output file + write_table( + direction_tel_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info( + "Processing and storing the subarray geometry prediction..." + ) + # Stack the telescope tables to the subarray table + direction_tel_tables = vstack(direction_tel_tables) + # Sort the table by the telescope event keys + direction_tel_tables.sort(TELESCOPE_EVENT_KEYS) + # Combine the telescope predictions to the subarray prediction using the stereo combiner + subarray_direction_table = self.geometry_stereo_combiner.predict_table( + direction_tel_tables + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if ( + subarray_direction_table[f"{self.prefix}_telescopes"].dtype + != np.bool_ + ): + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + (len(subarray_direction_table), len(self.dl1dh_reader.tel_ids)), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + subarray_direction_table[f"{self.prefix}_telescopes"] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] = True + # Overwrite the column with the boolean mask with fix length + subarray_direction_table[f"{self.prefix}_telescopes"] = ( + reco_telescopes + ) + # Save the prediction to the output file + write_table( + subarray_direction_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + ) + # Create the feature vector table if the DL1 features are enabled + if self.dl1_features: + self.log.info("Processing and storing dl1 feature vectors...") + feature_vector_table = super()._create_feature_vectors_table( + example_identifiers, + nonexample_identifiers, + classification_feature_vectors, + energy_feature_vectors, + direction_feature_vectors, + ) + # Loop over the selected telescopes and store the feature vectors + # for each telescope in the output file. The feature vectors are stored + # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = feature_vector_table["tel_id"] == tel_id + feature_vectors_tel_table = feature_vector_table[telescope_mask] + feature_vectors_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Save the prediction to the output file + write_table( + feature_vectors_tel_table, + self.output_path, + f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 feature vectors was stored in '%s' under '%s'", + self.output_path, + f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", + ) + + def _store_mc_telescope_pointing(self, all_identifiers): + """ + Store the telescope pointing table from MC simulation to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the telescope pointing information. + """ + # Create the pointing table for each telescope + pointing_info = [] + for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: + # Pointing table for the mono mode + tel_pointing = self.dl1dh_reader.get_tel_pointing(self.input_url, tel_id) + tel_pointing.rename_column("telescope_pointing_azimuth", "pointing_azimuth") + tel_pointing.rename_column( + "telescope_pointing_altitude", "pointing_altitude" + ) + # Join the prediction table with the telescope pointing table + tel_pointing = join( + left=tel_pointing, + right=all_identifiers, + keys=["obs_id", "tel_id"], + ) + # TODO: use keep_order for astropy v7.0.0 + tel_pointing.sort(TELESCOPE_EVENT_KEYS) + # Retrieve the example identifiers for the selected telescope + tel_pointing_table = Table( + { + "time": tel_pointing["time"], + "azimuth": tel_pointing["pointing_azimuth"], + "altitude": tel_pointing["pointing_altitude"], + } + ) + write_table( + tel_pointing_table, + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 telescope pointing table was stored in '%s' under '%s'", + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + ) + pointing_info.append(tel_pointing) + pointing_info = vstack(pointing_info) + return pointing_info + +def mono_tool(): + # Run the tool + mono_tool = MonoPredictCTLearnModel() + mono_tool.run() + +if __name__ == "main": + mono_tool() + +if __name__ == "__main__": + mono_tool() \ No newline at end of file diff --git a/ctlearn/tools/predict/predict_stereo.py b/ctlearn/tools/predict/predict_stereo.py new file mode 100644 index 00000000..c868fb70 --- /dev/null +++ b/ctlearn/tools/predict/predict_stereo.py @@ -0,0 +1,375 @@ +""" +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``. +""" + + +import numpy as np + +from astropy.table import ( + Table, + vstack, + join, + setdiff, +) + +from ctapipe.containers import ( + ParticleClassificationContainer, + ReconstructedGeometryContainer, + ReconstructedEnergyContainer, +) + +from ctapipe.io import read_table, write_table + +from ctapipe.reco.utils import add_defaults_and_meta +from dl1_data_handler.reader import ( + ProcessType, +) +from ctlearn.tools.predict.utils.predict_model import PredictCTLearnModel + + +SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" +FIXED_POINTING_GROUP = "/configuration/telescope/pointing" +POINTING_GROUP = "/dl1/monitoring/telescope/pointing" +SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" +DL1_TELESCOPE_GROUP = "/dl1/event/telescope" +DL1_SUBARRAY_GROUP = "/dl1/event/subarray" +DL2_SUBARRAY_GROUP = "/dl2/event/subarray" +DL2_TELESCOPE_GROUP = "/dl2/event/telescope" +SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] +TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] + +__all__ = ["StereoPredictCTLearnModel"] + +class StereoPredictCTLearnModel(PredictCTLearnModel): + """ + Tool to predict the gammaness, energy and arrival direction from R1/DL1 stereoscopic data using CTLearn models. + + This tool extends the ``PredictCTLearnModel`` to specifically handle stereoscopic R1/DL1 data. The prediction + is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. + It also stores the telescope/subarray pointing monitoring and DL1 feature vectors (if selected) in the output file. + + Attributes + ---------- + name : str + Name of the tool. + description : str + Description of the tool. + examples : str + Examples of how to use the tool. + + Methods + ------- + start() + Start the tool. + _store_mc_subarray_pointing(all_identifiers) + Store the subarray pointing table for the stereo mode for MC simulation. + """ + + name = "ctlearn-predict-stereo-model" + description = __doc__ + + examples = """ + To predict from pixel-wise image data in stereo mode using trained CTLearn models: + > ctlearn-predict-stereo-model \\ + --input_url input.dl1.h5 \\ + --PredictCTLearnModel.batch_size=16 \\ + --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ + --DLImageReader.channels=cleaned_image \\ + --DLImageReader.channels=cleaned_relative_peak_time \\ + --DLImageReader.image_mapper_type=BilinearMapper \\ + --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" \\ + --output output.dl2.h5 \\ + --PredictCTLearnModel.overwrite_tables=True \\ + """ + + def start(self): + self.log.info("Processing the telescope pointings...") + # Retrieve the IDs from the dl1dh for the prediction tables + example_identifiers = self.dl1dh_reader.unique_example_identifiers.copy() + example_identifiers.keep_columns(SUBARRAY_EVENT_KEYS) + all_identifiers = self.dl1dh_reader.subarray_trigger_table.copy() + all_identifiers.keep_columns(SUBARRAY_EVENT_KEYS + ["time"]) + nonexample_identifiers = setdiff( + all_identifiers, example_identifiers, keys=SUBARRAY_EVENT_KEYS + ) + nonexample_identifiers.remove_column("time") + # Construct the survival telescopes for each event of the example_identifiers + survival_telescopes = [] + for subarray_event in self.dl1dh_reader.example_identifiers_grouped.groups: + survival_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) + survival_tels = [ + self.dl1dh_reader.subarray.tel_indices[tel_id] + for tel_id in subarray_event["tel_id"].data + ] + survival_mask[survival_tels] = True + survival_telescopes.append(survival_mask) + # Add the survival telescopes to the example_identifiers + example_identifiers.add_column( + survival_telescopes, name=f"{self.prefix}_telescopes" + ) + # Pointing table for the stereo mode for MC simulation + if self.dl1dh_reader.process_type == ProcessType.Simulation: + pointing_info = self._store_mc_subarray_pointing(all_identifiers) + + # Pointing table for the observation mode + if self.dl1dh_reader.process_type == ProcessType.Observation: + pointing_info = super()._store_pointing(all_identifiers) + + self.log.info("Starting the prediction...") + classification_feature_vectors = None + if self.load_type_model_from is not None: + # Predict the energy of the primary particle + classification_table, classification_feature_vectors = ( + super()._predict_classification(example_identifiers) + ) + if self.dl2_subarray: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_prediction"], + shapes=[(len(nonexample_identifiers),)], + ) + classification_table = vstack([classification_table, nan_table]) + # Add is_valid column to the energy table + classification_table.add_column( + ~np.isnan( + classification_table[f"{self.prefix}_tel_prediction"].data, + dtype=bool, + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Rename the columns for the stereo mode + classification_table.rename_column( + f"{self.prefix}_tel_prediction", f"{self.prefix}_prediction" + ) + classification_table.rename_column( + f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" + ) + classification_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + classification_table, + ParticleClassificationContainer, + prefix=self.prefix, + ) + # Save the prediction to the output file + write_table( + classification_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + ) + energy_feature_vectors = None + if self.load_energy_model_from is not None: + # Predict the energy of the primary particle + energy_table, energy_feature_vectors = super()._predict_energy( + example_identifiers + ) + if self.dl2_subarray: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_energy"], + shapes=[(len(nonexample_identifiers),)], + ) + energy_table = vstack([energy_table, nan_table]) + # Add is_valid column to the energy table + energy_table.add_column( + ~np.isnan( + energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Rename the columns for the stereo mode + energy_table.rename_column( + f"{self.prefix}_tel_energy", f"{self.prefix}_energy" + ) + energy_table.rename_column( + f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" + ) + energy_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + energy_table, + ReconstructedEnergyContainer, + prefix=self.prefix, + ) + # Save the prediction to the output file + write_table( + energy_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + ) + direction_feature_vectors = None + if self.load_skydirection_model_from is not None: + # Join the prediction table with the telescope pointing table + example_identifiers = join( + left=example_identifiers, + right=pointing_info, + keys=SUBARRAY_EVENT_KEYS, + ) + # Predict the arrival direction of the primary particle + direction_table, direction_feature_vectors = super()._predict_skydirection( + example_identifiers + ) + if self.dl2_subarray: + # Transform the spherical coordinate offsets to sky coordinates + direction_table = super()._transform_spher_coord_offsets_to_sky( + direction_table + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_alt", f"{self.prefix}_az"], + shapes=[ + (len(nonexample_identifiers),), + (len(nonexample_identifiers),), + ], + ) + direction_table = vstack([direction_table, nan_table]) + # Add is_valid column to the direction table + direction_table.add_column( + ~np.isnan(direction_table[f"{self.prefix}_alt"].data, dtype=bool), + name=f"{self.prefix}_is_valid", + ) + direction_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + direction_table, + ReconstructedGeometryContainer, + prefix=self.prefix, + ) + # Save the prediction to the output file + write_table( + direction_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + ) + + # Create the feature vector table if the DL1 features are enabled + if self.dl1_features: + self.log.info("Processing and storing dl1 feature vectors...") + feature_vector_table = super()._create_feature_vectors_table( + example_identifiers, + nonexample_identifiers, + classification_feature_vectors, + energy_feature_vectors, + direction_feature_vectors, + ) + # Loop over the selected telescopes and store the feature vectors + # for each telescope in the output file. The feature vectors are stored + # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. + # Rename the columns for the stereo mode + feature_vector_table.rename_column( + f"{self.prefix}_tel_classification_feature_vectors", + f"{self.prefix}_classification_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefix}_tel_energy_feature_vectors", + f"{self.prefix}_energy_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefix}_tel_geometry_feature_vectors", + f"{self.prefix}_geometry_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" + ) + feature_vector_table.sort(SUBARRAY_EVENT_KEYS) + # Save the prediction to the output file + write_table( + feature_vector_table, + self.output_path, + f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 feature vectors was stored in '%s' under '%s'", + self.output_path, + f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", + ) + + def _store_mc_subarray_pointing(self, all_identifiers): + """ + Store the subarray pointing table from MC simulation to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the subarray pointing information. + """ + # Read the subarray pointing table + pointing_info = read_table( + self.input_url, + f"{SIMULATION_CONFIG_TABLE}", + ) + # Assuming min_az = max_az and min_alt = max_alt + pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) + pointing_info.rename_column("min_az", "pointing_azimuth") + pointing_info.rename_column("min_alt", "pointing_altitude") + # Join the prediction table with the telescope pointing table + pointing_info = join( + left=pointing_info, + right=all_identifiers, + keys=["obs_id"], + ) + # TODO: use keep_order for astropy v7.0.0 + pointing_info.sort(SUBARRAY_EVENT_KEYS) + # Create the pointing table + pointing_table = Table( + { + "time": pointing_info["time"], + "array_azimuth": pointing_info["pointing_azimuth"], + "array_altitude": pointing_info["pointing_altitude"], + "array_ra": np.nan * np.ones(len(pointing_info)), + "array_dec": np.nan * np.ones(len(pointing_info)), + } + ) + # Save the pointing table to the output file + write_table( + pointing_table, + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 subarray pointing table was stored in '%s' under '%s'", + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + ) + return pointing_info + +def stereo_tool(): + # Run the tool + stereo_tool = StereoPredictCTLearnModel() + stereo_tool.run() + +if __name__ == "main": + stereo_tool() +if __name__ == "__main__": + stereo_tool() \ No newline at end of file diff --git a/ctlearn/tools/predict/pytorch/predic_LST1_pytorch.py b/ctlearn/tools/predict/pytorch/predic_LST1_pytorch.py new file mode 100644 index 00000000..19b403c1 --- /dev/null +++ b/ctlearn/tools/predict/pytorch/predic_LST1_pytorch.py @@ -0,0 +1,312 @@ +""" +PyTorch prediction module for LST1 telescope data. +This module provides functionality to load trained models and perform predictions +on DL1 level data for particle type classification, energy estimation, and direction reconstruction. +""" + +from ctlearn.core.pytorch.net_utils import create_model, ModelHelper +import torch +from ctlearn.core.ctlearn_enum import Task, Mode +from ctapipe.io import read_table +from astropy.table import join +from dl1_data_handler.reader import get_unmapped_image +import numpy as np +from tqdm import tqdm +from pytorch_lightning.callbacks import Callback + +class GPUStatsLogger(Callback): + """ + PyTorch Lightning callback to log GPU memory statistics during training. + + This callback tracks GPU memory allocation and reservation at the end of each training epoch + and logs the statistics to TensorBoard. + """ + + def on_train_epoch_end(self, trainer, pl_module): + """ + Called at the end of each training epoch to log GPU memory statistics. + + Args: + trainer: PyTorch Lightning trainer instance + pl_module: The LightningModule being trained + """ + mem_allocated = torch.cuda.memory_allocated() + mem_reserved = torch.cuda.memory_reserved() + + trainer.logger.experiment.add_scalar( + "gpu_mem_allocated", mem_allocated, global_step=trainer.current_epoch + ) + trainer.logger.experiment.add_scalar( + "gpu_mem_reserved", mem_reserved, global_step=trainer.current_epoch + ) + + +def predictions(self): + """ + Perform predictions on input DL1 data using trained models. + + This function processes the input file in batches, applies quality cuts, + and generates predictions for particle type, energy, and/or direction + depending on the configured tasks. + + Returns: + tuple: Contains the following arrays: + - event_id: Event identifiers + - tel_azimuth: Telescope azimuth angles + - tel_altitude: Telescope altitude angles + - trigger_time: Event trigger times + - prediction: Particle type classification scores + - energy: Reconstructed energy values + - cam_coord_offset_x: Camera coordinate offset in x + - cam_coord_offset_y: Camera coordinate offset in y + - classification_fvs: Classification feature vectors + - energy_fvs: Energy estimation feature vectors + - direction_fvs: Direction reconstruction feature vectors + """ + + # Update channels based on log scaling config + if "normalization" in self.parameters and "apply_log_scaling" in self.parameters["normalization"]: + new_channels = list(self.channels) + if self.parameters["normalization"]["apply_log_scaling"][0] and not new_channels[0].startswith("log_"): + new_channels[0] = "log_" + new_channels[0] + self.channels = new_channels + + # Optimize batch size if requested + if self.optim_batch_size: + batch_size_found = False + batch = 256 + step = 16 + + while not batch_size_found: + from ctlearn.tools.predict.utils.optimaze_batch_size import test_batch + + # Load a test batch to find optimal batch size + dl1_table = read_table( + self.input_url, self.image_table_path, start=0, stop=batch + ) + dl1_table = join(left=dl1_table, right=self.parameter_table, keys=["event_id"]) + dl1_table = join(left=dl1_table, right=self.trigger_table, keys=["event_id"]) + + # Prepare test data + data = [] + for event in dl1_table: + image = get_unmapped_image(dl1_table[0], self.channels, self.transforms) + data.append(self.image_mapper.map_image(image)) + input_data = {"input": np.array(data)} + + imgs = input_data['input'][:, :, :, 0] + if len(self.channels) == 2: + peak_time = input_data['input'][:, :, :, 1] + + # Test batch size for each configured task + for task in self.tasks: + if task == Task.type: + batch_size_found = not test_batch( + self.type_model, + torch.tensor(imgs).unsqueeze(1).to(self.device), + torch.tensor(peak_time).unsqueeze(1).to(self.device), + self.device + ) + if task == Task.energy: + batch_size_found = not test_batch( + self.energy_model, + torch.tensor(imgs).unsqueeze(1).to(self.device), + torch.tensor(peak_time).unsqueeze(1).to(self.device), + self.device + ) + if task in [Task.cameradirection, Task.skydirection, Task.direction]: + batch_size_found = not test_batch( + self.cameradirection_model, + torch.tensor(imgs).unsqueeze(1).to(self.device), + torch.tensor(peak_time).unsqueeze(1).to(self.device), + self.device + ) + + batch += step + if not batch_size_found: + self.log.info(f"Batch size: {batch} OK") + + self.batch_size = batch - step + self.log.info(f"Optimized batch size: {self.batch_size}") + + # Initialize output arrays + event_id, tel_azimuth, tel_altitude, trigger_time = [], [], [], [] + prediction, energy, cam_coord_offset_x, cam_coord_offset_y = [], [], [], [] + classification_fvs, energy_fvs, direction_fvs = [], [], [] + + # Process input file in batches + for start in tqdm(range(0, self.table_length, self.batch_size), desc="Processing input file"): + stop = min(start + self.batch_size, self.table_length) + self.log.debug("Processing chunk from '%d' to '%d'.", start, stop - 1) + + # Read and join tables + dl1_table = read_table(self.input_url, self.image_table_path, start=start, stop=stop) + dl1_table = join(left=dl1_table, right=self.parameter_table, keys=["event_id"]) + dl1_table = join(left=dl1_table, right=self.trigger_table, keys=["event_id"]) + + # Apply quality selection + passes_quality_checks = np.ones(len(dl1_table), dtype=bool) + if self.quality_query: + passes_quality_checks = self.quality_query.get_table_mask(dl1_table) + dl1_table = dl1_table[passes_quality_checks] + + if len(dl1_table) == 0: + self.log.debug("No events passed the quality selection.") + continue + + # Prepare input data + data = [] + for event in dl1_table: + image = get_unmapped_image(event, self.channels, self.transforms) + data.append(self.image_mapper.map_image(image)) + input_data = {"input": np.array(data)} + + # Store metadata + event_id.extend(dl1_table["event_id"].data) + tel_azimuth.extend(dl1_table["tel_az"].data) + tel_altitude.extend(dl1_table["tel_alt"].data) + trigger_time.extend(dl1_table["time"].mjd) + + # Extract and clean image data + imgs = input_data['input'][:, :, :, 0] + if len(self.channels) == 2: + peak_time = input_data['input'][:, :, :, 1] + peak_time[peak_time < 0] = 0 + peak_time[np.isnan(peak_time)] = 0 + peak_time[np.isinf(peak_time)] = 0 + + imgs[imgs < 0] = 0 + imgs[np.isnan(imgs)] = 0 + imgs[np.isinf(imgs)] = 0 + + feature_vector = True + if self.parameters["normalization"]["apply_log_scaling"][0] == True: + imgs = imgs.astype(np.float32) + imgs = np.log10(imgs + 1.0) + if self.parameters["normalization"]["apply_log_scaling"][1] == True and len(self.channels) == 2: + peak_time = peak_time.astype(np.float32) + peak_time = np.log10(peak_time + 1.0) + + # Prepare unified input tensor + if len(self.channels) == 2: + input_tensor = torch.cat([ + torch.tensor(imgs).unsqueeze(1), + torch.tensor(peak_time).unsqueeze(1) + ], dim=1).to(self.device) + else: + input_tensor = torch.tensor(imgs).unsqueeze(1).to(self.device) + + # Run predictions for each configured task + for task in self.tasks: + if task == Task.type: + # Particle type classification + classification_pred, energy_pred, direction_pred = self.type_model(input_tensor) + + prediction.extend(torch.softmax(classification_pred[0], dim=1).cpu().detach().numpy()[:, 1]) + classification_fvs.extend(classification_pred[1].cpu().detach().numpy()) + + elif task == Task.energy: + # Energy estimation + classification_pred, energy_pred, direction_pred = self.energy_model(input_tensor) + + energy.extend(energy_pred[0].cpu().detach().numpy()) + if feature_vector: + energy_fvs.extend(energy_pred[1].cpu().detach().numpy()) + else: + energy_fvs.extend(np.array([[0]] * len(energy_pred[0]))) + + elif task in [Task.cameradirection, Task.skydirection, Task.direction]: + # Direction reconstruction + classification_pred, energy_pred, direction_pred = self.dirrection_model(input_tensor) + + cam_coord_offset_x.extend(direction_pred[0][:, 0].float().cpu().detach().numpy()) + cam_coord_offset_y.extend(direction_pred[0][:, 1].float().cpu().detach().numpy()) + if feature_vector: + direction_fvs.extend(direction_pred[1].cpu().detach().numpy()) + else: + direction_fvs.extend(np.array([[0]] * len(direction_pred[0]))) + + else: + raise ValueError( + f"task:{task.name} is not supported. Task must be type, direction or energy" + ) + + return (event_id, tel_azimuth, tel_altitude, trigger_time, prediction, energy, + cam_coord_offset_x, cam_coord_offset_y, classification_fvs, energy_fvs, direction_fvs) + + +def load_pytorch_model(self): + """ + Load PyTorch models from checkpoints for the configured tasks. + + This function creates and loads models for particle type classification, + energy estimation, and/or direction reconstruction based on the tasks + specified in the configuration. + + Returns: + torch.nn.Module: The last loaded model (for compatibility) + """ + model = None + from ctlearn.core.pytorch.model_collection import CTLearnPyTorchModel + + def load_pytorch_model_net(model_info, task_name, num_inputs, num_outputs): + model_name = model_info.get("model_name", "") + try: + component_cls = CTLearnPyTorchModel.non_abstract_subclasses().get(model_name) + if component_cls is not None: + params = model_info.get("parameters", {}).copy() + params.pop("task", None) + params.pop("num_inputs", None) + params.pop("num_outputs", None) + params["parent"] = self + component = component_cls( + task=task_name, + num_inputs=num_inputs, + num_outputs=num_outputs, + **params + ) + return component.model + except Exception as e: + self.log.warning(f"Failed to load model {model_name} as Component: {e}. Falling back to create_model.") + return create_model(model_info) + + num_inputs = 1 + + for task in self.tasks: + # Create model based on task type + if task == Task.type: + model_net = load_pytorch_model_net(self.parameters["model"]["model_type"], "type", num_inputs, 2) + check_point_path = self.parameters["data"]["type_checkpoint"] + + elif task == Task.energy: + model_net = load_pytorch_model_net(self.parameters["model"]["model_energy"], "energy", num_inputs, 1) + check_point_path = self.parameters["data"]["energy_checkpoint"] + + elif task in [Task.cameradirection, Task.skydirection, Task.direction]: + model_net = load_pytorch_model_net(self.parameters["model"]["model_direction"], "direction", num_inputs, 3) + check_point_path = self.parameters["data"]["direction_checkpoint"] + + else: + raise ValueError( + f"task:{task.name} is not supported. Task must be type, direction or energy" + ) + + # Load the model from checkpoint + model = ModelHelper.loadModel( + model_net, "", check_point_path, Mode.observation, device_str=self.device_str + ) + model.eval() + + # Assign model to appropriate attribute + if task == Task.type: + self.type_model = model + elif task == Task.energy: + self.energy_model = model + elif task in [Task.cameradirection, Task.skydirection, Task.direction]: + self.dirrection_model = model + else: + raise ValueError( + f"task:{task.name} is not supported. Task must be type, direction or energy" + ) + + return model diff --git a/ctlearn/tools/predict/pytorch/predic_model_pytorch.py b/ctlearn/tools/predict/pytorch/predic_model_pytorch.py new file mode 100644 index 00000000..0fcd89ca --- /dev/null +++ b/ctlearn/tools/predict/pytorch/predic_model_pytorch.py @@ -0,0 +1,144 @@ +""" +PyTorch model prediction module for CTLearn. +This module provides functionality to load trained models and perform batch predictions +on DL1 data for multiple tasks including particle classification, energy estimation, +and direction reconstruction. +""" + +from ctlearn.core.data_loader.loader import DLDataLoader +import torch +from tqdm import tqdm +import numpy as np +import inspect +from ctlearn.tools.predict.utils.load_model import load_model + + +def predict_with_model_pytorch(self, task): + """ + Load and predict with a CTLearn PyTorch model. + + This function loads a trained model from the specified path and performs predictions + on the provided data. It processes the data in batches and returns predictions for + particle type classification, energy estimation, and/or direction reconstruction + based on the configured task. + + Parameters + ---------- + task : Task + The task(s) to perform predictions for (type, energy, or direction). + + Returns + ------- + predict_data : dict + Dictionary containing prediction results with keys: + - 'type': Particle type classification probabilities (gammaness scores) + - 'energy': Reconstructed energy values + - 'cameradirection': Camera coordinate offsets for direction reconstruction + feature_vectors : None + Feature vectors (currently not extracted, placeholder for future implementation). + + Notes + ----- + The function automatically detects whether the model requires peak time information + by inspecting the model's forward method signature. Models can accept either one + input (image only) or two inputs (image and peak time). + """ + # Initialize batch size from configuration parameters + self.batch_size = self.parameters["hyp"]["batches"] + + # Create data loader for the specified task + # The DLDataLoader is initialized separately for each task to ensure robustness + channels = ["cleaned_image", "cleaned_peak_time"] + if self.parameters["normalization"]["apply_log_scaling"][0]: + channels[0] = "log_" + channels[0] + self.dl1dh_reader.channels = channels + data_loader = DLDataLoader.create( + framework="pytorch", + DLDataReader=self.dl1dh_reader, + indices=self.indices, + tasks=[task], + parameters=self.parameters, + use_augmentation=False, + batch_size=self.batch_size, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + + # Note: Handling of incomplete last batch + # In PyTorch, unlike Keras, we can process incomplete batches directly + # without needing a separate data loader. The code below is kept as reference + # for potential future use or compatibility with other frameworks. + + # 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.create( + # framework="pytorch", + # DLDataReader=self.dl1dh_reader, + # indices=last_batch_indices, + # tasks=task, + # parameters=self.parameters, + # use_augmentation=False, + # batch_size=self.last_batch_size, + # sort_by_intensity=self.sort_by_intensity, + # stack_telescope_images=self.stack_telescope_images, + # ) + + # Load the trained model from checkpoint + model = load_model(self) + + # Inspect model signature to determine number of inputs + # This allows the code to work with models that take either: + # - Single input: image only + # - Dual input: image and peak time + sig = inspect.signature(model.forward) + num_inputs = len(sig.parameters) + + # Initialize prediction data dictionary with empty lists + predict_data = {} + predict_data['type'] = [] + predict_data['energy'] = [] + predict_data["cameradirection"] = [] + + # Set model to evaluation mode (disables dropout, batch normalization, etc.) + model.eval() + + # Perform predictions without gradient computation (faster inference) + with torch.no_grad(): + for i, x in enumerate(tqdm(data_loader, desc="Processing", total=len(data_loader))): + # Skip empty batches + if len(x[0]['image']) == 0: + continue + + # Forward pass through the model + classification_pred, energy_pred, direction_pred = model( + x[0]['image'].to(self.device) + ) + + # Collect particle type classification predictions + if classification_pred[0] is not None: + # Apply softmax to get probability distribution and extract gammaness score + gammaness = torch.softmax(classification_pred[0], dim=1).cpu().detach().numpy() + predict_data['type'].extend(gammaness) + + # Collect energy estimation predictions + if energy_pred[0] is not None: + predict_data['energy'].extend(energy_pred[0].cpu().detach().numpy()) + + # Collect direction reconstruction predictions + if direction_pred[0] is not None: + predict_data["cameradirection"].extend(direction_pred[0].cpu().detach().numpy()) + + # Log progress every 100 batches + if i % 100 == 0: + self.log.info(f"Processed {i}/{len(data_loader)} events.") + + self.log.info("Processing completed.") + + # Convert lists to numpy arrays for efficient storage and further processing + predict_data["cameradirection"] = np.array(predict_data["cameradirection"]) + predict_data["type"] = np.array(predict_data["type"]) + predict_data["energy"] = np.array(predict_data["energy"]) + + # Return predictions and placeholder for feature vectors + return predict_data, None \ No newline at end of file diff --git a/ctlearn/tools/predict/utils/load_model.py b/ctlearn/tools/predict/utils/load_model.py new file mode 100644 index 00000000..3963142a --- /dev/null +++ b/ctlearn/tools/predict/utils/load_model.py @@ -0,0 +1,55 @@ +""" +Model loading utility module for CTLearn predictions. +This module provides a framework-agnostic interface for loading trained models, +supporting both Keras and PyTorch frameworks. +""" + + +def load_model(self): + """ + Load a trained model based on the configured framework type. + + This function acts as a dispatcher that delegates model loading to the appropriate + framework-specific implementation. It supports both Keras and PyTorch frameworks + and loads the model from the checkpoint path specified in the configuration. + + Parameters + ---------- + self : PredictionHandler + The prediction handler instance containing configuration parameters including: + - framework_type: str, either "keras" or "pytorch" + - Model checkpoint paths and other configuration parameters + + Returns + ------- + model : object + The loaded model ready for inference. Type depends on the framework: + - For Keras: keras.Model + - For PyTorch: torch.nn.Module + Returns None if the framework is not recognized. + + Raises + ------ + ImportError + If the specified framework's prediction module cannot be imported. + + Notes + ----- + The function automatically detects the framework type from the configuration + and imports the appropriate loading function dynamically to avoid unnecessary + dependencies when using only one framework. + """ + if self.framework_type == "keras": + # Load Keras model using framework-specific loader + from ctlearn.tools.predict.keras.predic_LST1_keras import load_keras_model + return load_keras_model(self) + + elif self.framework_type == "pytorch": + # Load PyTorch model using framework-specific loader + from ctlearn.tools.predict.pytorch.predic_LST1_pytorch import load_pytorch_model + return load_pytorch_model(self) + + else: + # Log error if framework is not recognized + self.log.error(f"Framework '{self.framework_type}' not found! Supported frameworks: 'keras', 'pytorch'") + return None \ No newline at end of file diff --git a/ctlearn/tools/predict/utils/optimaze_batch_size.py b/ctlearn/tools/predict/utils/optimaze_batch_size.py new file mode 100644 index 00000000..2cfc81cb --- /dev/null +++ b/ctlearn/tools/predict/utils/optimaze_batch_size.py @@ -0,0 +1,174 @@ +""" +Batch size optimization utilities for CTLearn predictions. +This module provides functions to test and find optimal batch sizes for model inference, +helping to maximize GPU utilization while avoiding out-of-memory errors. +""" + +import torch + + +def test_batch(model, imgs, peak_time, device): + """ + Test if a model can process a given batch without memory errors. + + This function tests whether a pre-prepared batch of images and peak times + can be successfully processed by the model without encountering out-of-memory + (OOM) errors. It's used to validate batch sizes during optimization. + + Parameters + ---------- + model : torch.nn.Module + The PyTorch model to test. + imgs : torch.Tensor or array-like + Batch of images to process. Will be converted to tensor if necessary. + peak_time : torch.Tensor or array-like + Batch of peak time information. Will be converted to tensor if necessary. + device : torch.device or str + Device to run the test on (e.g., 'cuda:0' or 'cpu'). + + Returns + ------- + bool + True if the batch can be processed successfully, False if OOM error occurs. + + Raises + ------ + RuntimeError + If a RuntimeError other than OOM occurs during processing. + + Notes + ----- + The function automatically clears the CUDA cache after each test to ensure + clean memory state for subsequent tests. + """ + # Move model to specified device and set to evaluation mode + model.to(device) + model.eval() + + # Ensure inputs are tensors and move to device + if not torch.is_tensor(imgs): + imgs = torch.as_tensor(imgs).to(device) + else: + imgs = imgs.to(device) + + if not torch.is_tensor(peak_time): + peak_time = torch.as_tensor(peak_time).to(device) + else: + peak_time = peak_time.to(device) + + try: + # Attempt forward pass without gradient computation + with torch.no_grad(): + _ = model(imgs, peak_time) + + # Clear CUDA cache to free memory + torch.cuda.empty_cache() + return True + + except RuntimeError as e: + # Check if error is due to out of memory + if "out of memory" in str(e).lower(): + torch.cuda.empty_cache() + return False + else: + # Re-raise unexpected errors after cleaning up + torch.cuda.empty_cache() + raise e + + +def find_max_batch_size(self, model, imgs, peak_time, device, start_bs=8, step=8, max_bs=512): + """ + Find the maximum batch size that can be processed without OOM errors. + + This function performs a binary-like search to find the largest batch size + that can be successfully processed by the model on the given device. It starts + with a small batch size and incrementally increases until an OOM error occurs. + + Parameters + ---------- + self : object + Reference to the parent object (for potential logging or configuration access). + model : torch.nn.Module + The PyTorch model to test. + imgs : torch.Tensor or array-like + Sample images to use for testing. Only the first image is used and replicated. + peak_time : torch.Tensor or array-like + Sample peak time data. Only the first value is used and replicated. + device : torch.device or str + Device to run tests on (e.g., 'cuda:0' or 'cpu'). + start_bs : int, optional + Initial batch size to start testing with. Default is 8. + step : int, optional + Increment step for batch size increases. Default is 8. + max_bs : int, optional + Maximum batch size to test. Default is 512. + + Returns + ------- + int + The maximum batch size that can be processed without OOM errors. + + Raises + ------ + RuntimeError + If a RuntimeError other than OOM occurs during testing. + + Notes + ----- + - The function replicates a single image/peak_time to create test batches + - CUDA cache is cleared after each test to ensure accurate memory measurements + - Progress is printed to console with emoji indicators for status + """ + batch_size = start_bs + + # Move model to device and set to evaluation mode + model.to(device) + model.eval() + + # Ensure inputs are tensors + if not torch.is_tensor(imgs): + imgs = torch.as_tensor(imgs) + if not torch.is_tensor(peak_time): + peak_time = torch.as_tensor(peak_time) + + # Disable gradient computation for efficiency + with torch.no_grad(): + while batch_size <= max_bs: + try: + # Prepare image batch + batch_imgs = imgs[:1] # Take first image as template + if batch_imgs.ndim == 3: + batch_imgs = batch_imgs.unsqueeze(1) # Add channel dimension if needed + # Replicate to create batch of desired size + batch_imgs = batch_imgs.repeat(batch_size, 1, 1, 1).to(device) + + # Prepare peak_time batch + value = peak_time[:1, 0, 0] # Extract representative value + batch_peaks = value.unsqueeze(1) # Shape [1, 1] + batch_peaks = batch_peaks.repeat(batch_size, 1) # Replicate for batch + batch_peaks = batch_peaks.unsqueeze(-1).unsqueeze(-1).to(device) # Shape [batch, 1, 1, 1] + + # Attempt forward pass + _ = model(batch_imgs, batch_peaks) + + # Clean up tensors and cache + del batch_imgs, batch_peaks, _ + torch.cuda.empty_cache() + + print(f"✅ Batch size {batch_size} OK") + batch_size += step + + except RuntimeError as e: + if "out of memory" in str(e).lower(): + # OOM encountered, return previous successful batch size + print(f"💥 OOM at batch size {batch_size}") + torch.cuda.empty_cache() + return batch_size - step + else: + # Unexpected error, clean up and re-raise + print(f"❌ Unexpected error at batch size {batch_size}: {e}") + torch.cuda.empty_cache() + raise e + + # If we reached max_bs without OOM, return it (minus step to be safe) + return batch_size - step diff --git a/ctlearn/tools/predict/utils/predict_model.py b/ctlearn/tools/predict/utils/predict_model.py new file mode 100644 index 00000000..2ca79a25 --- /dev/null +++ b/ctlearn/tools/predict/utils/predict_model.py @@ -0,0 +1,1196 @@ +""" +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``. +""" + +import atexit +import pathlib +import numpy as np +import os +import tensorflow as tf +import keras +import threading +from ctlearn.core.ctlearn_enum import Task + +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, +) +from ctlearn.tools.train.pytorch.utils import ( + sanity_check, + read_configuration, + expected_structure, +) + +from ctapipe.containers import ( + ParticleClassificationContainer, + ReconstructedGeometryContainer, + ReconstructedEnergyContainer, +) +from ctapipe.coordinates import CameraFrame, NominalFrame +from ctapipe.core import Tool +from ctapipe.core.tool import ToolConfigurationError +from ctapipe.core.traits import ( + Bool, + Float, + Int, + Path, + flag, + Set, + Dict, + List, + CaselessStrEnum, + ComponentName, + Unicode, + classes_with_traits, +) +from ctapipe.monitoring.interpolation import PointingInterpolator +from ctapipe.io import read_table, write_table, HDF5Merger +from ctapipe.reco.reconstructor import ReconstructionProperty +from ctapipe.reco.stereo_combination import StereoCombiner +from ctapipe.reco.utils import add_defaults_and_meta +from dl1_data_handler.reader import ( + DLDataReader, + ProcessType, + LST_EPOCH, +) +from ctlearn.core.data_loader.loader import DLDataLoader +from ctlearn.utils import monitor_progress + +SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" +FIXED_POINTING_GROUP = "/configuration/telescope/pointing" +POINTING_GROUP = "/dl1/monitoring/telescope/pointing" +SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" +DL1_TELESCOPE_GROUP = "/dl1/event/telescope" +DL1_SUBARRAY_GROUP = "/dl1/event/subarray" +DL2_SUBARRAY_GROUP = "/dl2/event/subarray" +DL2_TELESCOPE_GROUP = "/dl2/event/telescope" +SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] +TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] + +__all__ = ["PredictCTLearnModel"] + +class PredictCTLearnModel(Tool): + """ + Base tool to predict the gammaness, energy and arrival direction from R1/DL1 data using CTLearn models. + + 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``. + 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. + + Attributes + ---------- + input_url : pathlib.Path + Input ctapipe HDF5 files including pixel-wise image or waveform data. + use_HDF5Merger : bool + Set whether to use the HDF5Merger component to copy the selected tables from the input file to the output file. + dl1_features : bool + Set whether to include the dl1 feature vectors in the output file. + dl2_telescope : bool + Set whether to include dl2 telescope-event-wise data in the output file. + dl2_subarray : bool + Set whether to include dl2 subarray-event-wise data in the output file. + dl1dh_reader : dl1_data_handler.reader.DLDataReader + DLDataReader object to read the data. + dl1dh_reader_type : str + Type of the DLDataReader to use for the prediction. + stack_telescope_images : bool + Set whether to stack the telescope images in the data loader. Requires ``stereo``. + sort_by_intensity : bool + Set whether to sort the telescope images by intensity in the data loader. Requires ``stereo``. + 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) or directory (Keras2) for the classification of the primary particle type. + load_energy_model_from : pathlib.Path + Path to a Keras model file (Keras3) or directory (Keras2) for the regression of the primary particle energy. + load_cameradirection_model_from : pathlib.Path + Path to a Keras model file (Keras3) or directory (Keras2) for the regression + of the primary particle arrival direction based on camera coordinate offsets. + load_cameradirection_model_from : pathlib.Path + Path to a Keras model file (Keras3) or directory (Keras2) 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. + overwrite_tables : bool + Overwrite the table in the output file if it exists. + keras_verbose : int + 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. + indices : list of int + List of indices for the data loaders. + batch_size : int + Size of the batch to perform inference of the neural network. + last_batch_size : int + Size of the last batch in the data loaders. + + Methods + ------- + setup() + Set up the tool. + finish() + Finish the tool. + _predict_with_model(model_path) + Load and predict with a CTLearn model. + _predict_classification(example_identifiers) + Predict the classification of the primary particle type. + _predict_energy(example_identifiers) + Predict the energy of the primary particle. + _predict_cameradirection(example_identifiers) + Predict the arrival direction of the primary particle based on camera coordinate offsets. + _predict_skydirection(example_identifiers) + Predict the arrival direction of the primary particle based on spherical coordinate offsets. + _transform_cam_coord_offsets_to_sky(table) + Transform to camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + _transform_spher_coord_offsets_to_sky(table) + Transform to spherical coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + _create_nan_table(nonexample_identifiers, columns, shapes) + Create a table with NaNs for missing predictions. + _store_pointing(all_identifiers) + Store the telescope pointing table from to the output file. + _create_feature_vectors_table(example_identifiers, nonexample_identifiers, classification_feature_vectors, energy_feature_vectors, direction_feature_vectors) + Create the table for the DL1 feature vectors. + """ + + input_url = Path( + help="Input ctapipe HDF5 files including pixel-wise image or waveform data", + allow_none=True, + exists=True, + directory_ok=False, + file_ok=True, + ).tag(config=True) + + use_HDF5Merger = Bool( + default_value=True, + allow_none=False, + help=( + "Set whether to use the HDF5Merger component to copy the selected tables " + "from the input file to the output file. CAUTION: This can only be used " + "if the output file not exists." + ), + ).tag(config=True) + + dl1_features = Bool( + default_value=False, + allow_none=False, + help="Set whether to include the dl1 feature vectors in the output file.", + ).tag(config=True) + + dl2_telescope = Bool( + default_value=True, + allow_none=False, + help="Set whether to include dl2 telescope-event-wise data in the output file.", + ).tag(config=True) + + dl2_subarray = Bool( + default_value=True, + allow_none=False, + help="Set whether to include dl2 subarray-event-wise data in the output file.", + ).tag(config=True) + + dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( + config=True + ) + + stack_telescope_images = Bool( + default_value=False, + allow_none=False, + help=( + "Set whether to stack the telescope images in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + sort_by_intensity = Bool( + default_value=False, + allow_none=False, + help=( + "Set whether to sort the telescope images by intensity in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + prefix = Unicode( + default_value="CTLearn", + allow_none=False, + help="Name of the reconstruction algorithm used to generate the dl2 data.", + ).tag(config=True) + + load_type_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the classification " + "of the primary particle type." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_energy_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " + "of the primary particle energy." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_cameradirection_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " + "of the primary particle arrival direction based on camera coordinate offsets." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_skydirection_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " + "of the primary particle arrival direction based on spherical coordinate offsets." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + batch_size = Int( + default_value=64, + allow_none=False, + help="Size of the batch to perform inference of the neural network.", + ).tag(config=True) + + output_path = Path( + default_value="./output.dl2.h5", + allow_none=False, + help="Output path to save the dl2 prediction results", + ).tag(config=True) + + overwrite_tables = Bool( + default_value=True, + allow_none=False, + help="Overwrite the table in the output file if it exists", + ).tag(config=True) + + pytorch_config_file = Path( + default_value=None, + allow_none=True, + help="Pytorch config file", + ).tag(config=True) + + # Unified Hardware Architecture & Execution Strategy + device = CaselessStrEnum( + ["cuda", "cpu", "mps"], + default_value="cuda", + help="Device to use: cuda, cpu, or mps", + ).tag(config=True) + + devices = List( + trait=Int(), + default_value=[0], + help="List of GPU device IDs to use", + ).tag(config=True) + + strategy = Unicode( + default_value="auto", + help="Multi-GPU strategy", + ).tag(config=True) + + # Event Cut-offs + leakage_intensity_cutoff = Float( + default_value=0.2, + help="Events with leakage intensity greater than this value are removed", + ).tag(config=True) + + intensity_cutoff = Float( + default_value=50.0, + help="Events with Hillas intensity below this value are removed", + ).tag(config=True) + + # Normalizations + apply_log_scaling = List( + trait=Bool(), + default_value=[True, True], + help="List specifying whether to apply log10(X+1.0) scaling for [charge, peak_time]", + ).tag(config=True) + + use_clean = Bool( + default_value=True, + help="Use the image with the applied mask", + ).tag(config=True) + + use_clean_dvr = Bool( + default_value=False, + help="Use clean DVR mask", + ).tag(config=True) + + type_mu = Float( + default_value=0.0, + help="Mean for type channel normalization", + ).tag(config=True) + + type_sigma = Float( + default_value=1000.0, + help="Std dev for type channel normalization", + ).tag(config=True) + + dir_mu = Float( + default_value=0.0, + help="Mean for direction channel normalization", + ).tag(config=True) + + dir_sigma = Float( + default_value=1000.0, + help="Std dev for direction channel normalization", + ).tag(config=True) + + energy_mu = Float( + default_value=0.0, + help="Mean for energy channel normalization", + ).tag(config=True) + + energy_sigma = Float( + default_value=1000.0, + help="Std dev for energy channel normalization", + ).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) + + framework_type = CaselessStrEnum( + ["pytorch", "keras"], + default_value="keras", + help="Framework to use: pytorch or keras", + ).tag(config=True) + + aliases = { + ("i", "input_url"): "PredictCTLearnModel.input_url", + ("t", "type_model"): "PredictCTLearnModel.load_type_model_from", + ("e", "energy_model"): "PredictCTLearnModel.load_energy_model_from", + ( + "d", + "cameradirection_model", + ): "PredictCTLearnModel.load_cameradirection_model_from", + ("s", "skydirection_model"): "PredictCTLearnModel.load_skydirection_model_from", + ("o", "output"): "PredictCTLearnModel.output_path", + ("f", "framework"): "PredictCTLearnModel.framework_type", + ("p", "pytorch_config_file"): "PredictCTLearnModel.pytorch_config_file", + "device": "PredictCTLearnModel.device", + "devices": "PredictCTLearnModel.devices", + "strategy": "PredictCTLearnModel.strategy", + "type_mu": "PredictCTLearnModel.type_mu", + "type_sigma": "PredictCTLearnModel.type_sigma", + "dir_mu": "PredictCTLearnModel.dir_mu", + "dir_sigma": "PredictCTLearnModel.dir_sigma", + "energy_mu": "PredictCTLearnModel.energy_mu", + "energy_sigma": "PredictCTLearnModel.energy_sigma", + } + + flags = { + **flag( + "dl1-features", + "PredictCTLearnModel.dl1_features", + "Include dl1 features", + "Exclude dl1 features", + ), + **flag( + "dl2-telescope", + "PredictCTLearnModel.dl2_telescope", + "Include dl2 telescope-event-wise data in the output file", + "Exclude dl2 telescope-event-wise data in the output file", + ), + **flag( + "dl2-subarray", + "PredictCTLearnModel.dl2_subarray", + "Include dl2 telescope-event-wise data in the output file", + "Exclude dl2 telescope-event-wise data in the output file", + ), + **flag( + "use-HDF5Merger", + "PredictCTLearnModel.use_HDF5Merger", + "Copy data using the HDF5Merger component (CAUTION: This can not be used if the output file already exists)", + "Do not copy data using the HDF5Merger component", + ), + **flag( + "r0-waveforms", + "HDF5Merger.r0_waveforms", + "Include r0 waveforms", + "Exclude r0 waveforms", + ), + **flag( + "r1-waveforms", + "HDF5Merger.r1_waveforms", + "Include r1 waveforms", + "Exclude r1 waveforms", + ), + **flag( + "dl1-parameters", + "HDF5Merger.dl1_parameters", + "Include dl1 parameters", + "Exclude dl1 parameters", + ), + **flag( + "dl1-images", + "HDF5Merger.dl1_images", + "Include dl1 images", + "Exclude dl1 images", + ), + **flag( + "true-parameters", + "HDF5Merger.true_parameters", + "Include true parameters", + "Exclude true parameters", + ), + **flag( + "true-images", + "HDF5Merger.true_images", + "Include true images", + "Exclude true images", + ), + } + + classes = classes_with_traits(DLDataReader) + + def setup(self): + if self.framework_type == "pytorch": + import torch + if self.pytorch_config_file is not None: + self.log.info(f"Using {self.pytorch_config_file} config file for pytorch framework") + legacy_params = read_configuration(self.pytorch_config_file) + sanity_check(legacy_params, expected_structure) + + def get_conf_val(key1, key2, trait_name, default_val): + in_config = False + if "PredictCTLearnModel" in self.config and trait_name in self.config["PredictCTLearnModel"]: + in_config = True + if not in_config: + return legacy_params.get(key1, {}).get(key2, default_val) + return getattr(self, trait_name) + + self.device_str = get_conf_val("arch", "device", "device", self.device.name if hasattr(self.device, "name") else str(self.device)) + self.type_mu = get_conf_val("normalization", "type_mu", "type_mu", self.type_mu) + self.type_sigma = get_conf_val("normalization", "type_sigma", "type_sigma", self.type_sigma) + self.dir_mu = get_conf_val("normalization", "dir_mu", "dir_mu", self.dir_mu) + self.dir_sigma = get_conf_val("normalization", "dir_sigma", "dir_sigma", self.dir_sigma) + self.energy_mu = get_conf_val("normalization", "energy_mu", "energy_mu", self.energy_mu) + self.energy_sigma = get_conf_val("normalization", "energy_sigma", "energy_sigma", self.energy_sigma) + self.leakage_intensity_cutoff = get_conf_val("cut-off", "leakage_intensity", "leakage_intensity_cutoff", self.leakage_intensity_cutoff) + self.intensity_cutoff = get_conf_val("cut-off", "intensity", "intensity_cutoff", self.intensity_cutoff) + + self.parameters = legacy_params + else: + self.log.info("No legacy config file provided. Using standard Traitlets configuration for PyTorch.") + self.device_str = self.device.name if hasattr(self.device, "name") else str(self.device) + + self.parameters = { + "data": { + "type_checkpoint": self.load_type_model_from, + "energy_checkpoint": self.load_energy_model_from, + "direction_checkpoint": self.load_cameradirection_model_from or self.load_skydirection_model_from, + }, + "hyp": { + "batches": self.batch_size, + "dynamic_batches": True, + }, + "cut-off": { + "leakage_intensity": self.leakage_intensity_cutoff, + "intensity": self.intensity_cutoff, + }, + "normalization": { + "apply_log_scaling": self.apply_log_scaling, + "use_clean": self.use_clean, + "use_clean_dvr": self.use_clean_dvr, + "type_mu": self.type_mu, + "type_sigma": self.type_sigma, + "dir_mu": self.dir_mu, + "dir_sigma": self.dir_sigma, + "energy_mu": self.energy_mu, + "energy_sigma": self.energy_sigma, + }, + "arch": { + "device": self.device_str, + "devices": self.devices, + "strategy": self.strategy, + }, + "model": { + "model_type": { + "model_name": "DoubleBBEfficientNet", + "parameters": { + "model_variant": "efficientnet-b3", + "task": "type", + "num_outputs": 2, + "device_str": self.device_str, + "energy_bins": None, + } + }, + "model_energy": { + "model_name": "ThinResNet", + "parameters": { + "task": "energy", + "num_inputs": 1, + "num_outputs": 1, + "num_blocks": [3, 4, 6, 3], + "dropout": 0.1, + "use_bn": False, + } + }, + "model_direction": { + "model_name": "ThinResNet_DBB", + "parameters": { + "task": "direction", + "num_inputs": 1, + "num_outputs": 3, + "num_blocks": [3, 4, 6, 3], + "dropout": 0.1, + "use_bn": False, + } + } + } + } + self.device = torch.device(self.device_str) + self.tasks = [] + if self.load_type_model_from is not None: + self.tasks.append(Task.type) + if self.load_energy_model_from is not None: + self.tasks.append(Task.energy) + if self.load_cameradirection_model_from is not None or self.load_skydirection_model_from is not None: + self.tasks.append(Task.direction) + + # Check if the ctapipe HDF5Merger component is enabled + if os.path.exists(self.output_path): + self.log.warning( + "The output file '{self.output_path}' already exists. Disabling HDF5Merger the flag '--no-use-HDF5Merger' will be not use." + ) + self.use_HDF5Merger = False + + if self.use_HDF5Merger: + if os.path.exists(self.output_path): + raise ToolConfigurationError( + f"The output file '{self.output_path}' already exists. Please use " + "'--no-use-HDF5Merger' to disable the usage of the HDF5Merger component." + ) + # Copy selected tables from the input file to the output file + self.log.info("Copying to output destination.") + stop_event = threading.Event() + monitor_thread = threading.Thread(target=monitor_progress, args=(self.input_url, self.output_path, stop_event, self.log)) + monitor_thread.start() + + try: + with HDF5Merger(self.output_path, parent=self) as merger: + merger(self.input_url) + finally: + stop_event.set() + monitor_thread.join() + + else: + self.log.info( + "No copy to output destination, since the usage of the HDF5Merger component is disabled." + ) + + # 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) + + # Set up the data reader + self.log.info("Loading data reader:") + self.log.info("For a large dataset, this may take a while...") + self.dl1dh_reader = DLDataReader.from_name( + self.dl1dh_reader_type, + input_url_signal=[self.input_url], + parent=self, + ) + self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) + # Check if the number of events is enough to form a batch + if self.dl1dh_reader._get_n_events() < self.batch_size: + raise ToolConfigurationError( + f"{self.dl1dh_reader._get_n_events()} events are not enough " + f"to form a batch of size {self.batch_size}. Reduce the batch size." + ) + # 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 + ) + + def finish(self): + self.log.info("Tool is shutting down") + + def _predict_with_model(self, model_path, task): + """ + Load and predict with a CTLearn 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. + + Parameters + ---------- + model_path : str + Path to a Keras model file (Keras3) or directory (Keras2). + + Returns + ------- + predict_data : astropy.table.Table + Table containing the prediction results. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + predict_data = None + feature_vectors = None + + if self.framework_type == "keras": + from ctlearn.tools.predict.keras.predic_model_keras import predict_with_model + predict_data, feature_vectors = predict_with_model(self,model_path) + return predict_data, feature_vectors + + if self.framework_type == "pytorch": + from ctlearn.tools.predict.pytorch.predic_model_pytorch import predict_with_model_pytorch + + task = None + if model_path == self.load_type_model_from: + task = Task.type + elif model_path == self.load_energy_model_from: + task = Task.energy + elif model_path in [self.load_cameradirection_model_from, self.load_skydirection_model_from]: + task = Task.direction + else: + task = Task.type + + predict_data, feature_vectors = predict_with_model_pytorch(self, task) + return predict_data, feature_vectors + return predict_data, feature_vectors + + def _predict_classification(self, example_identifiers): + """ + Predict the classification of the primary particle type. + + This method uses a pre-trained type model to predict the type of the primary particle + for a given set of example identifiers. The predicted classification score ('gammaness') + is added to the example identifiers table. + + Parameters: + ----------- + classification_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + predicted classification score ('gammaness'). + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the classification of the primary particle type..." + ) + # Predict the data using the loaded type_model + predict_data, feature_vectors = self._predict_with_model( + self.load_type_model_from, Task.type + ) + # Create prediction table and add the predicted classification score ('gammaness') + classification_table = example_identifiers.copy() + classification_table.add_column( + predict_data["type"].T[1], name=f"{self.prefix}_tel_prediction" + ) + return classification_table, feature_vectors + + def _predict_energy(self, example_identifiers): + """ + Predict the energy of the primary particle. + + This method uses a pre-trained energy model to predict the energy of the primary particle + for a given set of example identifiers. The predicted energy is then converted from + log10(TeV) to TeV and added to the example identifiers table. + + Parameters: + ----------- + energy_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed energy in TeV. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info("Predicting for the regression of the primary particle energy...") + # Predict the data using the loaded energy_model + predict_data, feature_vectors = self._predict_with_model( + self.load_energy_model_from, Task.energy + ) + # Convert the reconstructed energy from log10(TeV) to TeV + reco_energy = u.Quantity( + np.power(10, np.squeeze(predict_data["energy"])), + unit=u.TeV, + ) + print(reco_energy) + print(reco_energy.shape) + print(example_identifiers) + # Create prediction table and add the reconstructed energy in TeV + energy_table = example_identifiers.copy() + energy_table.add_column(reco_energy, name=f"{self.prefix}_tel_energy") + return energy_table, feature_vectors + + def _predict_cameradirection(self, example_identifiers): + """ + Predict the arrival direction of the primary particle based on camera coordinate offsets. + + This method uses a pre-trained direction model to predict the arrival direction of the + primary particle for a given set of example identifiers. The predicted camera coordinate offsets + is added to the example identifiers table. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + + Returns: + -------- + cameradirection_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed camera coordinate offsets in x and y. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the regression of the primary particle arrival direction based on camera coordinate offsets..." + ) + # Predict the data using the loaded direction_model + predict_data, feature_vectors = self._predict_with_model( + self.load_cameradirection_model_from, Task.direction + ) + # For the direction task, the prediction is the camera coordinate offset in x and y + # from the telescope pointing. + cam_coord_offset_x = u.Quantity(predict_data["cameradirection"].T[0], unit=u.m) + cam_coord_offset_y = u.Quantity(predict_data["cameradirection"].T[1], unit=u.m) + # Create prediction table and add the reconstructed energy in TeV + cameradirection_table = example_identifiers.copy() + cameradirection_table.add_column(cam_coord_offset_x, name="cam_coord_offset_x") + cameradirection_table.add_column(cam_coord_offset_y, name="cam_coord_offset_y") + return cameradirection_table, feature_vectors + + def _predict_skydirection(self, example_identifiers): + """ + Predict the arrival direction of the primary particle based on spherical coordinate offsets. + + This method uses a pre-trained direction model to predict the arrival direction of the primary + particle for a given set of example identifiers. The predicted spherical coordinate offsets is + added to the example identifiers table. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + + Returns: + -------- + skydirection_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed spherical coordinate offsets in fov_lon and fov_lat. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the regression of the primary particle arrival direction based on spherical coordinate offsets..." + ) + # Predict the data using the loaded direction_model + predict_data, feature_vectors = self._predict_with_model( + self.load_skydirection_model_from + ) + # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat + # from the telescope pointing. + fov_lon = u.Quantity(predict_data["skydirection"].T[0], unit=u.deg) + fov_lat = u.Quantity(predict_data["skydirection"].T[1], unit=u.deg) + # Create prediction table and add the reconstructed energy in TeV + skydirection_table = example_identifiers.copy() + skydirection_table.add_column(fov_lon, name="fov_lon") + skydirection_table.add_column(fov_lat, name="fov_lat") + return skydirection_table, feature_vectors + + def _transform_cam_coord_offsets_to_sky(self, table) -> Table: + """ + Transform the predicted camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + + This method converts the predicted camera coordinate offsets w.r.t. the telescope pointing + in the provided table to Alt/Az coordinates. It also removes the unnecessary columns + from the table that do not the ctapipe DL2 data format. + + Parameters: + ----------- + table : astropy.table.Table + A Table containing the trigger time, telescope pointing, and predicted camera coordinate offsets. + + Returns: + -------- + table : astropy.table.Table + A Table with the Alt/Az coordinates following the ctapipe DL2 data format. + """ + # Get the telescope ID from the table + tel_id = table["tel_id"][0] + # Set the telescope position + tel_ground_frame = self.dl1dh_reader.subarray.tel_coords[ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] + # Set the trigger timestamp based on the process type + if self.dl1dh_reader.process_type == ProcessType.Simulation: + trigger_time = LST_EPOCH + elif self.dl1dh_reader.process_type == ProcessType.Observation: + trigger_time = table["time"] + # Set the telescope pointing with the trigger timestamp and the telescope position + altaz = AltAz( + location=tel_ground_frame.to_earth_location(), + obstime=trigger_time, + ) + # Set the telescope pointing + tel_pointing = SkyCoord( + az=table["pointing_azimuth"], + alt=table["pointing_altitude"], + frame=altaz, + ) + # Set the camera frame with the focal length and rotation of the camera + camera_frame = CameraFrame( + focal_length=self.dl1dh_reader.subarray.tel[ + tel_id + ].camera.geometry.frame.focal_length, + rotation=self.dl1dh_reader.pix_rotation[tel_id], + telescope_pointing=tel_pointing, + ) + # Set the camera coordinate offset + cam_coord_offset = SkyCoord( + x=table["cam_coord_offset_x"], + y=table["cam_coord_offset_y"], + frame=camera_frame, + ) + # tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] + # Transform the true Alt/Az coordinates to camera coordinates + reco_direction = cam_coord_offset.transform_to(altaz) + # Add the reconstructed direction (az, alt) to the prediction table + table.add_column(reco_direction.az.to(u.deg), name=f"{self.prefix}_tel_az") + table.add_column(reco_direction.alt.to(u.deg), name=f"{self.prefix}_tel_alt") + # Remove unnecessary columns from the table that do not the ctapipe DL2 data format + table.remove_columns( + [ + "time", + "pointing_azimuth", + "pointing_altitude", + "cam_coord_offset_x", + "cam_coord_offset_y", + ] + ) + return table + + def _transform_spher_coord_offsets_to_sky(self, table) -> Table: + """ + Transform the predicted spherical offsets w.r.t. the telescope pointing to Alt/Az coordinates. + + This method converts the predicted spherical offsets w.r.t. the telescope pointing + in the provided table to Alt/Az coordinates. It also removes the unnecessary columns + from the table that do not the ctapipe DL2 data format. + + Parameters: + ----------- + table : astropy.table.Table + A Table containing the trigger time, telescope pointing, and predicted spherical offsets. + + Returns: + -------- + table : astropy.table.Table + A Table with the Alt/Az coordinates following the ctapipe DL2 data format. + """ + + # Set the trigger timestamp based on the process type + if self.dl1dh_reader.process_type == ProcessType.Simulation: + trigger_time = LST_EPOCH + elif self.dl1dh_reader.process_type == ProcessType.Observation: + trigger_time = table["time"] + # Set the AltAz frame with the reference location and time + altaz = AltAz( + location=self.dl1dh_reader.subarray.reference_location, + obstime=trigger_time, + ) + # Set the array pointing + array_pointing = SkyCoord( + az=table["pointing_azimuth"], + alt=table["pointing_altitude"], + frame=altaz, + ) + # Set the nominal frame with the array pointing + nom_frame = NominalFrame( + origin=array_pointing, + location=self.dl1dh_reader.subarray.reference_location, + obstime=trigger_time, + ) + # Set the reco direction in (fov_lon, fov_lat) coordinates + reco_direction = SkyCoord( + fov_lon=table["fov_lon"], + fov_lat=table["fov_lat"], + frame=nom_frame, + ) + # Transform the reco direction from nominal frame to the AltAz frame + sky_coord = reco_direction.transform_to(altaz) + # Add the reconstructed direction (az, alt) to the prediction table + table.add_column(sky_coord.az.to(u.deg), name=f"{self.prefix}_az") + table.add_column(sky_coord.alt.to(u.deg), name=f"{self.prefix}_alt") + # Remove unnecessary columns from the table that do not the ctapipe DL2 data format + table.remove_columns( + [ + "time", + "pointing_azimuth", + "pointing_altitude", + "fov_lon", + "fov_lat", + ] + ) + return table + + def _create_nan_table(self, nonexample_identifiers, columns, shapes): + """ + Create a table with NaNs for missing predictions. + + This method creates a table with NaNs for missing predictions for the non-example identifiers. + In stereo mode, the table also a column for the valid telescopes is added with all False values. + + Parameters: + ----------- + nonexample_identifiers : astropy.table.Table + Table containing the non-example identifiers. + columns : list of str + List of column names to create in the table. + shapes : list of shapes + List of shapes for the columns to create in the table. + + Returns: + -------- + nan_table : astropy.table.Table + Table containing NaNs for missing predictions. + """ + # Create a table with NaNs for missing predictions + nan_table = nonexample_identifiers.copy() + for column_name, shape in zip(columns, shapes): + nan_table.add_column(np.full(shape, np.nan), name=column_name) + # Add that no telescope is valid for the non-example identifiers in stereo mode + if self.dl1dh_reader.mode == "stereo": + nan_table.add_column( + np.zeros( + (len(nonexample_identifiers), len(self.dl1dh_reader.tel_ids)), + dtype=bool, + ), + name=f"{self.prefix}_telescopes", + ) + return nan_table + + def _store_pointing(self, all_identifiers): + """ + Store the telescope pointing table from to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the telescope pointing information. + """ + + # Initialize the pointing interpolator from ctapipe + pointing_interpolator = PointingInterpolator( + bounds_error=False, extrapolate=True + ) + pointing_info = [] + for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: + # Get the telescope pointing from the dl1dh reader + tel_pointing = self.dl1dh_reader.telescope_pointings[f"tel_{tel_id:03d}"] + # Add the telescope pointing table to the pointing interpolator + pointing_interpolator.add_table(tel_id, tel_pointing) + tel_identifiers = all_identifiers.copy() + if self.dl1dh_reader.mode == "mono": + tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] + # Interpolate the telescope pointing + tel_altitude, tel_azimuth = pointing_interpolator( + tel_id, tel_identifiers["time"] + ) + tel_identifiers.add_column(tel_azimuth, name="pointing_azimuth") + tel_identifiers.add_column(tel_altitude, name="pointing_altitude") + pointing_info.append(tel_identifiers) + if self.dl1dh_reader.mode == "mono": + tel_pointing_table = Table( + { + "time": tel_identifiers["time"], + "azimuth": tel_identifiers["pointing_azimuth"], + "altitude": tel_identifiers["pointing_altitude"], + } + ) + write_table( + tel_pointing_table, + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 telescope pointing table was stored in '%s' under '%s'", + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + ) + pointing_info = vstack(pointing_info) + if self.dl1dh_reader.mode == "stereo": + # Group the pointing information by subarray event keys + # TODO: This needs to be debugged with SST1M data + pointing_info_grouped = pointing_info.group_by(SUBARRAY_EVENT_KEYS) + pointing_mean = pointing_info_grouped.groups.aggregate(np.mean) + pointing_info = join( + all_identifiers, + pointing_mean, + keys=SUBARRAY_EVENT_KEYS, + ) + # TODO: use keep_order for astropy v7.0.0 + pointing_info.sort(SUBARRAY_EVENT_KEYS) + # Create the pointing table + pointing_table = Table( + { + "time": pointing_info["time"], + "array_azimuth": pointing_info["pointing_azimuth"], + "array_altitude": pointing_info["pointing_altitude"], + "array_ra": np.nan * np.ones(len(pointing_info)), + "array_dec": np.nan * np.ones(len(pointing_info)), + } + ) + # Save the pointing table to the output file + write_table( + pointing_table, + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 subarray pointing table was stored in '%s' under '%s'", + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + ) + return pointing_info + + def _create_feature_vectors_table( + self, + example_identifiers, + nonexample_identifiers=None, + classification_feature_vectors=None, + energy_feature_vectors=None, + direction_feature_vectors=None, + ): + """ + Create the table for the DL1 feature vectors. + + This method creates a table with the DL1 feature vectors for the example identifiers and fill NaNs for + non-example identifiers. The feature vectors are stored in the columns of the table. The table also + contains a column for the valid predictions. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + nonexample_identifiers : astropy.table.Table or None + Table containing the non-example identifiers to fill the NaNs. + classification_feature_vectors : np.ndarray or None + Array containing the classification feature vectors. + energy_feature_vectors : np.ndarray or None + Array containing the energy feature vectors. + direction_feature_vectors : np.ndarray or None + Array containing the direction feature vectors. + + Returns: + -------- + feature_vector_table : astropy.table.Table + Table containing the DL1 feature vectors for the example and non-example identifiers. + """ + # Create the feature vector table + feature_vector_table = example_identifiers.copy() + feature_vector_table.remove_columns( + ["pointing_azimuth", "pointing_altitude", "time"] + ) + columns_list, shapes_list = [], [] + if classification_feature_vectors is not None: + is_valid_col = ~np.isnan( + np.min(classification_feature_vectors, axis=1), dtype=bool + ) + feature_vector_table.add_column( + classification_feature_vectors, + name=f"{self.prefix}_tel_classification_feature_vectors", + ) + if nonexample_identifiers is not None: + columns_list.append(f"{self.prefix}_tel_classification_feature_vectors") + shapes_list.append( + ( + len(nonexample_identifiers), + classification_feature_vectors.shape[1], + ) + ) + if energy_feature_vectors is not None: + is_valid_col = ~np.isnan(np.min(energy_feature_vectors, axis=1), dtype=bool) + feature_vector_table.add_column( + energy_feature_vectors, name=f"{self.prefix}_tel_energy_feature_vectors" + ) + if nonexample_identifiers is not None: + columns_list.append(f"{self.prefix}_tel_energy_feature_vectors") + shapes_list.append( + ( + len(nonexample_identifiers), + energy_feature_vectors.shape[1], + ) + ) + if direction_feature_vectors is not None: + is_valid_col = ~np.isnan( + np.min(direction_feature_vectors, axis=1), dtype=bool + ) + feature_vector_table.add_column( + direction_feature_vectors, + name=f"{self.prefix}_tel_geometry_feature_vectors", + ) + if nonexample_identifiers is not None: + columns_list.append(f"{self.prefix}_tel_geometry_feature_vectors") + shapes_list.append( + ( + len(nonexample_identifiers), + direction_feature_vectors.shape[1], + ) + ) + # Produce output table with NaNs for missing predictions + if nonexample_identifiers is not None: + if len(nonexample_identifiers) > 0: + nan_table = self._create_nan_table( + nonexample_identifiers, + columns=columns_list, + shapes=shapes_list, + ) + feature_vector_table = vstack([feature_vector_table, nan_table]) + is_valid_col = np.concatenate( + (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) + ) + # Add is_valid column to the feature vector table + feature_vector_table.add_column( + is_valid_col, + name=f"{self.prefix}_tel_is_valid", + ) + return feature_vector_table + + +if __name__ == "__main__": + pass \ No newline at end of file diff --git a/ctlearn/tools/predict_LST1.py b/ctlearn/tools/predict_LST1.py index 58069e56..e814f63d 100644 --- a/ctlearn/tools/predict_LST1.py +++ b/ctlearn/tools/predict_LST1.py @@ -4,7 +4,10 @@ import numpy as np import tables -import keras +try: + import keras +except ImportError: + keras = None from astropy import units as u from astropy.coordinates import AltAz, SkyCoord from astropy.table import Table, join, setdiff, vstack @@ -16,6 +19,7 @@ ReconstructedEnergyContainer, ) from ctapipe.coordinates import CameraFrame +from ctapipe.coordinates import CameraFrame from ctapipe.core import Tool from ctapipe.core.tool import ToolConfigurationError from ctapipe.core.traits import ( @@ -28,28 +32,50 @@ ComponentName, Dict, UseEnum, + UseEnum, classes_with_traits, ) from ctapipe.instrument.optics import FocalLengthKind +from ctapipe.instrument.optics import FocalLengthKind from ctapipe.io import read_table, write_table -from ctapipe.io.hdf5dataformat import ( - DL1_SUBARRAY_TRIGGER_TABLE, - DL1_TEL_GROUP, - DL1_TEL_PARAMETERS_GROUP, - DL1_TEL_POINTING_GROUP, - DL1_TEL_TRIGGER_TABLE, - DL2_TEL_PARTICLETYPE_GROUP, - DL2_TEL_ENERGY_GROUP, - DL2_TEL_GEOMETRY_GROUP, - DL2_SUBARRAY_PARTICLETYPE_GROUP, - DL2_SUBARRAY_ENERGY_GROUP, - DL2_SUBARRAY_GEOMETRY_GROUP, -) + +DL0_TEL_POINTING_GROUP = "/dl0/event/telescope/pointing" +DL1_SUBARRAY_GROUP = "/dl1/event/subarray" +DL1_SUBARRAY_POINTING_GROUP = "/dl1/event/subarray/pointing" +DL1_SUBARRAY_TRIGGER_TABLE = "/dl1/event/subarray/trigger" +DL1_TEL_GROUP = "/dl1/event/telescope" +DL1_TEL_CALIBRATION_GROUP = "/dl1/event/telescope/calibration" +DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP = "/dl1/event/telescope/illuminator_throughput" +DL1_TEL_IMAGES_GROUP = "/dl1/event/telescope/image" +DL1_TEL_MUON_GROUP = "/dl1/event/telescope/muon" +DL1_TEL_MUON_THROUGHPUT_GROUP = "/dl1/event/telescope/muon_throughput" +DL1_TEL_OPTICAL_PSF_GROUP = "/dl1/event/telescope/optical_psf" +DL1_TEL_PARAMETERS_GROUP = "/dl1/event/telescope/parameters" +DL1_TEL_POINTING_GROUP = "/dl1/event/telescope/pointing" +DL1_TEL_TRIGGER_TABLE = "/dl1/event/telescope/trigger" +DL2_EVENT_STATISTICS_GROUP = "/dl2/event/subarray/statistics" +FIXED_POINTING_GROUP = "/configuration/telescope/pointing" +R0_TEL_GROUP = "/r0/event/telescope" +R1_TEL_GROUP = "/r1/event/telescope" +SIMULATION_IMAGES_GROUP = "/simulation/event/telescope/images" +SIMULATION_IMPACT_GROUP = "/simulation/event/telescope/impact" +SIMULATION_PARAMETERS_GROUP = "/simulation/event/telescope/parameters" +SIMULATION_RUN_TABLE = "/simulation/run_config" +SIMULATION_SHOWER_TABLE = "/simulation/event/subarray/shower" +DL2_TEL_PARTICLETYPE_GROUP = "/dl2/event/telescope/classification" +DL2_TEL_ENERGY_GROUP = "/dl2/event/telescope/energy" +DL2_TEL_GEOMETRY_GROUP = "/dl2/event/telescope/geometry" +DL2_SUBARRAY_GROUP = "/dl2/event/subarray" +DL2_SUBARRAY_PARTICLETYPE_GROUP = "/dl2/event/subarray/classification" +DL2_SUBARRAY_ENERGY_GROUP = "/dl2/event/subarray/energy" +DL2_SUBARRAY_GEOMETRY_GROUP = "/dl2/event/subarray/geometry" + from ctapipe.reco.utils import add_defaults_and_meta +from ctlearn.core.keras.model import LoadedModel from ctlearn import __version__ as ctlearn_version -from ctlearn.utils import get_lst1_subarray_description, validate_trait_dict - +from ctlearn.utils import get_lst1_subarray_description +from ctlearn.utils import validate_trait_dict from dl1_data_handler.image_mapper import ImageMapper from dl1_data_handler.reader import ( get_unmapped_image, @@ -184,11 +210,11 @@ class LST1PredictionTool(Tool): "cleaned_image", "peak_time", "relative_peak_time", - "cleaned_peak_time", + "cleaned_peak_time", "cleaned_relative_peak_time", ] ), - default_value=["cleaned_image", "cleaned_relative_peak_time"], + default_value=["cleaned_image", "cleaned_peak_time"], allow_none=False, help=( "Set the input channels to be loaded from the DL1 event data. " @@ -243,6 +269,8 @@ class LST1PredictionTool(Tool): ("e", "energy_model"): "LST1PredictionTool.load_energy_model_from", ("d", "cameradirection_model"): "LST1PredictionTool.load_cameradirection_model_from", ("o", "output"): "LST1PredictionTool.output_path", + ("ch", "channels"): "LST1PredictionTool.channels", + } flags = { @@ -274,7 +302,12 @@ class LST1PredictionTool(Tool): } - classes = classes_with_traits(ImageMapper) + @property + def classes(self): + return [ + type(self), + LST1PredictionTool, + ] + classes_with_traits(ImageMapper) def setup(self): self.log.info("ctlearn version %s", ctlearn_version) @@ -290,6 +323,9 @@ def setup(self): # 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)) + + if keras is None: + raise ImportError("TensorFlow/Keras is required for prediction. Install it with 'pip install ctlearn[tf]' or 'pip install ctlearn[all]'.") # Load the models from the specified paths if self.load_type_model_from is not None: @@ -952,4 +988,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..be6f1180 100644 --- a/ctlearn/tools/predict_model.py +++ b/ctlearn/tools/predict_model.py @@ -8,8 +8,12 @@ import numpy as np import tables -import tensorflow as tf -import keras +try: + import tensorflow as tf + import keras +except ImportError: + tf = None + keras = None from astropy import units as u from astropy.coordinates.earth import EarthLocation @@ -39,44 +43,45 @@ flag, Dict, ComponentName, + CaselessStrEnum, classes_with_traits, ) from ctapipe.monitoring.interpolation import PointingInterpolator from ctapipe.instrument import SubarrayDescription from ctapipe.io import read_table, write_table, HDF5Merger from ctapipe.io.datalevels import DataLevel -from ctapipe.io.hdf5dataformat import ( - DL0_TEL_POINTING_GROUP, - DL1_SUBARRAY_GROUP, - DL1_SUBARRAY_POINTING_GROUP, - DL1_SUBARRAY_TRIGGER_TABLE, - DL1_TEL_GROUP, - DL1_TEL_CALIBRATION_GROUP, - DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP, - DL1_TEL_IMAGES_GROUP, - DL1_TEL_MUON_GROUP, - DL1_TEL_MUON_THROUGHPUT_GROUP, - DL1_TEL_OPTICAL_PSF_GROUP, - DL1_TEL_PARAMETERS_GROUP, - DL1_TEL_POINTING_GROUP, - DL1_TEL_TRIGGER_TABLE, - DL2_EVENT_STATISTICS_GROUP, - FIXED_POINTING_GROUP, - R0_TEL_GROUP, - R1_TEL_GROUP, - SIMULATION_IMAGES_GROUP, - SIMULATION_IMPACT_GROUP, - SIMULATION_PARAMETERS_GROUP, - SIMULATION_RUN_TABLE, - SIMULATION_SHOWER_TABLE, - DL2_TEL_PARTICLETYPE_GROUP, - DL2_TEL_ENERGY_GROUP, - DL2_TEL_GEOMETRY_GROUP, - DL2_SUBARRAY_GROUP, - DL2_SUBARRAY_PARTICLETYPE_GROUP, - DL2_SUBARRAY_ENERGY_GROUP, - DL2_SUBARRAY_GEOMETRY_GROUP, -) + +DL0_TEL_POINTING_GROUP = "/dl0/event/telescope/pointing" +DL1_SUBARRAY_GROUP = "/dl1/event/subarray" +DL1_SUBARRAY_POINTING_GROUP = "/dl1/event/subarray/pointing" +DL1_SUBARRAY_TRIGGER_TABLE = "/dl1/event/subarray/trigger" +DL1_TEL_GROUP = "/dl1/event/telescope" +DL1_TEL_CALIBRATION_GROUP = "/dl1/event/telescope/calibration" +DL1_TEL_ILLUMINATOR_THROUGHPUT_GROUP = "/dl1/event/telescope/illuminator_throughput" +DL1_TEL_IMAGES_GROUP = "/dl1/event/telescope/image" +DL1_TEL_MUON_GROUP = "/dl1/event/telescope/muon" +DL1_TEL_MUON_THROUGHPUT_GROUP = "/dl1/event/telescope/muon_throughput" +DL1_TEL_OPTICAL_PSF_GROUP = "/dl1/event/telescope/optical_psf" +DL1_TEL_PARAMETERS_GROUP = "/dl1/event/telescope/parameters" +DL1_TEL_POINTING_GROUP = "/dl1/event/telescope/pointing" +DL1_TEL_TRIGGER_TABLE = "/dl1/event/telescope/trigger" +DL2_EVENT_STATISTICS_GROUP = "/dl2/event/subarray/statistics" +FIXED_POINTING_GROUP = "/configuration/telescope/pointing" +R0_TEL_GROUP = "/r0/event/telescope" +R1_TEL_GROUP = "/r1/event/telescope" +SIMULATION_IMAGES_GROUP = "/simulation/event/telescope/images" +SIMULATION_IMPACT_GROUP = "/simulation/event/telescope/impact" +SIMULATION_PARAMETERS_GROUP = "/simulation/event/telescope/parameters" +SIMULATION_RUN_TABLE = "/simulation/run_config" +SIMULATION_SHOWER_TABLE = "/simulation/event/subarray/shower" +DL2_TEL_PARTICLETYPE_GROUP = "/dl2/event/telescope/classification" +DL2_TEL_ENERGY_GROUP = "/dl2/event/telescope/energy" +DL2_TEL_GEOMETRY_GROUP = "/dl2/event/telescope/geometry" +DL2_SUBARRAY_GROUP = "/dl2/event/subarray" +DL2_SUBARRAY_PARTICLETYPE_GROUP = "/dl2/event/subarray/classification" +DL2_SUBARRAY_ENERGY_GROUP = "/dl2/event/subarray/energy" +DL2_SUBARRAY_GEOMETRY_GROUP = "/dl2/event/subarray/geometry" + from ctapipe.reco.reconstructor import ReconstructionProperty from ctapipe.reco.stereo_combination import StereoCombiner from ctapipe.reco.utils import add_defaults_and_meta @@ -86,7 +91,7 @@ LST_EPOCH, ) from ctlearn import __version__ as ctlearn_version -from ctlearn.core.loader import DLDataLoader +from ctlearn.core.data_loader.loader import DLDataLoader from ctlearn.utils import validate_trait_dict # Convienient constants for column names and table keys @@ -342,6 +347,12 @@ class PredictCTLearnModel(Tool): ), ).tag(config=True) + framework_type = CaselessStrEnum( + ["pytorch", "keras"], + default_value="keras", + help="Framework to use: pytorch or keras", + ).tag(config=True) + aliases = { ("i", "input_url"): "PredictCTLearnModel.input_url", ("t", "type_model"): "PredictCTLearnModel.load_type_model_from", @@ -352,6 +363,7 @@ class PredictCTLearnModel(Tool): ): "PredictCTLearnModel.load_cameradirection_model_from", ("s", "skydirection_model"): "PredictCTLearnModel.load_skydirection_model_from", ("o", "output"): "PredictCTLearnModel.output_path", + ("f", "framework"): "PredictCTLearnModel.framework_type", } flags = { @@ -417,7 +429,13 @@ class PredictCTLearnModel(Tool): ), } - classes = classes_with_traits(DLDataReader) + @property + def classes(self): + return [ + type(self), + PredictCTLearnModel, + HDF5Merger, + ] + classes_with_traits(DLDataReader) def setup(self): self.activity_start_time = Time.now() @@ -432,6 +450,10 @@ def setup(self): self.output_path, dl2_subarray=False, dl2_telescope=False, parent=self ) as merger: merger(self.input_url) + + if tf is None: + raise ImportError("TensorFlow is required for prediction. Install it with 'pip install ctlearn[tf]' or 'pip install ctlearn[all]'.") + # Create a MirroredStrategy. self.strategy = tf.distribute.MirroredStrategy() atexit.register(self.strategy._extended._collective_ops._lock.locked) # type: ignore @@ -675,7 +697,7 @@ def deduplicate_first_valid( return unique(t, keys=list(keys), keep="first") - def _predict_with_model(self, model_path): + def _predict_with_model(self, model_path, task): """ Load and predict with a CTLearn model. @@ -686,6 +708,8 @@ def _predict_with_model(self, model_path): ---------- model_path : str Path to a Keras model file (Keras3). + task : str + The task for which prediction is being made. Returns ------- @@ -694,124 +718,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( - self.dl1dh_reader, - self.indices, - tasks=[], - batch_size=self.batch_size * self.strategy.num_replicas_in_sync, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - # Keras is only considering the last complete batch. - # In prediction mode we don't want to loose the last - # uncomplete batch, so we are creating an additional - # batch generator for the remaining events. - 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( - self.dl1dh_reader, - last_batch_indices, - tasks=[], - batch_size=self.last_batch_size, - sort_by_intensity=self.sort_by_intensity, - stack_telescope_images=self.stack_telescope_images, - ) - # Load the model from the specified path - model = keras.saving.load_model(model_path) - prediction_colname = ( - "type" - if isinstance(model.layers[-1], keras.layers.Softmax) - else model.layers[-1].name - ) - backbone_model, feature_vectors = None, None - if self.dl1_features: - # Get the backbone model which is the second layer of the model - backbone_model = 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) - # 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 - ) - except ValueError as err: - if str(err).startswith("Input 0 of layer"): - raise ToolConfigurationError( - "Model input shape does not match the prediction data. " - "This is usually caused by selecting the wrong telescope_id. " - "Please ensure the telescope configuration matches the one used for training." - ) from err - raise - # Apply the head model with the feature vectors to retrieve the prediction - predict_data = Table( - { - prediction_colname: head.predict( - feature_vectors, verbose=self.keras_verbose - ) - } - ) - # 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 - ) - feature_vectors = np.concatenate( - (feature_vectors, feature_vectors_last_batch) - ) - predict_data = vstack( - [ - predict_data, - Table( - { - prediction_colname: head.predict( - feature_vectors_last_batch, - verbose=self.keras_verbose, - ) - } - ), - ] - ) - else: - # Predict the data using the loaded model - try: - predict_data = model.predict(data_loader, verbose=self.keras_verbose) - except ValueError as err: - if str(err).startswith("Input 0 of layer"): - raise ToolConfigurationError( - "Model input shape does not match the prediction data. " - "This is usually caused by selecting the wrong telescope_id. " - "Please ensure the telescope configuration matches the one used for training." - ) from err - raise - # Create a astropy table with the prediction results - # The classification task has a softmax layer as the last layer - # which returns the probabilities for each class in an array, while - # the regression tasks have output neurons which returns the - # predicted value for the task in a dictionary. - if prediction_colname == "type": - predict_data = Table({prediction_colname: predict_data}) - else: - predict_data = Table(predict_data) - # 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 - ) - if model.layers[-1].name == "type": - predict_data_last_batch = Table( - {prediction_colname: predict_data_last_batch} - ) - else: - predict_data_last_batch = Table(predict_data_last_batch) - predict_data = vstack([predict_data, predict_data_last_batch]) + if self.framework_type == "keras": + from ctlearn.tools.predict.keras.predic_model_keras import predict_with_model + predict_data, feature_vectors = predict_with_model(self, model_path, task) + + return predict_data, feature_vectors + + return predict_data, feature_vectors def _predict_particletype(self, example_identifiers): @@ -835,7 +748,7 @@ def _predict_particletype(self, example_identifiers): ) # Predict the data using the loaded type_model predict_data, feature_vectors = self._predict_with_model( - self.load_type_model_from + self.load_type_model_from, "type" ) # Create prediction table and add the predicted classification score ('gammaness') particletype_table = example_identifiers.copy() @@ -863,7 +776,7 @@ def _predict_energy(self, example_identifiers): self.log.info("Predicting for the regression of the primary particle energy...") # Predict the data using the loaded energy_model predict_data, feature_vectors = self._predict_with_model( - self.load_energy_model_from + self.load_energy_model_from, "energy" ) # Convert the reconstructed energy from log10(TeV) to TeV reco_energy = u.Quantity( @@ -903,7 +816,7 @@ def _predict_cameradirection(self, example_identifiers): ) # Predict the data using the loaded direction_model predict_data, feature_vectors = self._predict_with_model( - self.load_cameradirection_model_from + self.load_cameradirection_model_from, "cameradirection" ) # For the direction task, the prediction is the camera coordinate offset in x and y # from the telescope pointing. @@ -941,7 +854,7 @@ def _predict_skydirection(self, example_identifiers): ) # Predict the data using the loaded direction_model predict_data, feature_vectors = self._predict_with_model( - self.load_skydirection_model_from + self.load_skydirection_model_from, "skydirection" ) # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat # from the telescope pointing. @@ -2310,10 +2223,14 @@ def _store_mc_subarray_pointing(self, all_identifiers): Table containing the subarray pointing information. """ # Read the subarray pointing table - pointing_info = read_table( - self.input_url, - SIMULATION_RUN_TABLE, - ) + try: + pointing_info = read_table( + self.input_url, + SIMULATION_RUN_TABLE, + ) + except Exception as e: + self.log.warning("Could not read simulation run table: %s", e) + return None # Assuming min_az = max_az and min_alt = max_alt pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) pointing_info.rename_column("min_az", "pointing_azimuth") diff --git a/ctlearn/tools/predict_model_main.py b/ctlearn/tools/predict_model_main.py new file mode 100644 index 00000000..4b64001b --- /dev/null +++ b/ctlearn/tools/predict_model_main.py @@ -0,0 +1,1884 @@ +""" +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``. +""" + +import atexit +import pathlib +import numpy as np +import os +try: + import tensorflow as tf + import keras +except ImportError: + tf = None + keras = None + +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, +) + +from ctapipe.containers import ( + ParticleClassificationContainer, + ReconstructedGeometryContainer, + ReconstructedEnergyContainer, +) +from ctapipe.coordinates import CameraFrame, NominalFrame +from ctapipe.core import Tool +from ctapipe.core.tool import ToolConfigurationError +from ctapipe.core.traits import ( + Bool, + Int, + Path, + flag, + Set, + Dict, + List, + CaselessStrEnum, + ComponentName, + Unicode, + classes_with_traits, +) +from ctapipe.monitoring.interpolation import PointingInterpolator +from ctapipe.io import read_table, write_table, HDF5Merger +from ctapipe.reco.reconstructor import ReconstructionProperty +from ctapipe.reco.stereo_combination import StereoCombiner +from ctapipe.reco.utils import add_defaults_and_meta +from dl1_data_handler.reader import ( + DLDataReader, + ProcessType, + LST_EPOCH, +) +from ctlearn.core.loader import DLDataLoader + +SIMULATION_CONFIG_TABLE = "/configuration/simulation/run" +FIXED_POINTING_GROUP = "/configuration/telescope/pointing" +POINTING_GROUP = "/dl1/monitoring/telescope/pointing" +SUBARRAY_POINTING_GROUP = "/dl1/monitoring/subarray/pointing" +DL1_TELESCOPE_GROUP = "/dl1/event/telescope" +DL1_SUBARRAY_GROUP = "/dl1/event/subarray" +DL2_SUBARRAY_GROUP = "/dl2/event/subarray" +DL2_TELESCOPE_GROUP = "/dl2/event/telescope" +SUBARRAY_EVENT_KEYS = ["obs_id", "event_id"] +TELESCOPE_EVENT_KEYS = ["obs_id", "event_id", "tel_id"] + +__all__ = [ + "PredictCTLearnModel", + "MonoPredictCTLearnModel", + "StereoPredictCTLearnModel", +] + + +class PredictCTLearnModel(Tool): + """ + Base tool to predict the gammaness, energy and arrival direction from R1/DL1 data using CTLearn models. + + 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``. + 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. + + Attributes + ---------- + input_url : pathlib.Path + Input ctapipe HDF5 files including pixel-wise image or waveform data. + use_HDF5Merger : bool + Set whether to use the HDF5Merger component to copy the selected tables from the input file to the output file. + dl1_features : bool + Set whether to include the dl1 feature vectors in the output file. + dl2_telescope : bool + Set whether to include dl2 telescope-event-wise data in the output file. + dl2_subarray : bool + Set whether to include dl2 subarray-event-wise data in the output file. + dl1dh_reader : dl1_data_handler.reader.DLDataReader + DLDataReader object to read the data. + dl1dh_reader_type : str + Type of the DLDataReader to use for the prediction. + stack_telescope_images : bool + Set whether to stack the telescope images in the data loader. Requires ``stereo``. + sort_by_intensity : bool + Set whether to sort the telescope images by intensity in the data loader. Requires ``stereo``. + 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) or directory (Keras2) for the classification of the primary particle type. + load_energy_model_from : pathlib.Path + Path to a Keras model file (Keras3) or directory (Keras2) for the regression of the primary particle energy. + load_cameradirection_model_from : pathlib.Path + Path to a Keras model file (Keras3) or directory (Keras2) for the regression + of the primary particle arrival direction based on camera coordinate offsets. + load_cameradirection_model_from : pathlib.Path + Path to a Keras model file (Keras3) or directory (Keras2) 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. + overwrite_tables : bool + Overwrite the table in the output file if it exists. + keras_verbose : int + 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. + indices : list of int + List of indices for the data loaders. + batch_size : int + Size of the batch to perform inference of the neural network. + last_batch_size : int + Size of the last batch in the data loaders. + + Methods + ------- + setup() + Set up the tool. + finish() + Finish the tool. + _predict_with_model(model_path) + Load and predict with a CTLearn model. + _predict_classification(example_identifiers) + Predict the classification of the primary particle type. + _predict_energy(example_identifiers) + Predict the energy of the primary particle. + _predict_cameradirection(example_identifiers) + Predict the arrival direction of the primary particle based on camera coordinate offsets. + _predict_skydirection(example_identifiers) + Predict the arrival direction of the primary particle based on spherical coordinate offsets. + _transform_cam_coord_offsets_to_sky(table) + Transform to camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + _transform_spher_coord_offsets_to_sky(table) + Transform to spherical coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + _create_nan_table(nonexample_identifiers, columns, shapes) + Create a table with NaNs for missing predictions. + _store_pointing(all_identifiers) + Store the telescope pointing table from to the output file. + _create_feature_vectors_table(example_identifiers, nonexample_identifiers, classification_feature_vectors, energy_feature_vectors, direction_feature_vectors) + Create the table for the DL1 feature vectors. + """ + + input_url = Path( + help="Input ctapipe HDF5 files including pixel-wise image or waveform data", + allow_none=True, + exists=True, + directory_ok=False, + file_ok=True, + ).tag(config=True) + + use_HDF5Merger = Bool( + default_value=True, + allow_none=False, + help=( + "Set whether to use the HDF5Merger component to copy the selected tables " + "from the input file to the output file. CAUTION: This can only be used " + "if the output file not exists." + ), + ).tag(config=True) + + dl1_features = Bool( + default_value=False, + allow_none=False, + help="Set whether to include the dl1 feature vectors in the output file.", + ).tag(config=True) + + dl2_telescope = Bool( + default_value=True, + allow_none=False, + help="Set whether to include dl2 telescope-event-wise data in the output file.", + ).tag(config=True) + + dl2_subarray = Bool( + default_value=True, + allow_none=False, + help="Set whether to include dl2 subarray-event-wise data in the output file.", + ).tag(config=True) + + dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( + config=True + ) + + stack_telescope_images = Bool( + default_value=False, + allow_none=False, + help=( + "Set whether to stack the telescope images in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + sort_by_intensity = Bool( + default_value=False, + allow_none=False, + help=( + "Set whether to sort the telescope images by intensity in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + prefix = Unicode( + default_value="CTLearn", + allow_none=False, + help="Name of the reconstruction algorithm used to generate the dl2 data.", + ).tag(config=True) + + load_type_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the classification " + "of the primary particle type." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_energy_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " + "of the primary particle energy." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_cameradirection_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " + "of the primary particle arrival direction based on camera coordinate offsets." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + load_skydirection_model_from = Path( + default_value=None, + help=( + "Path to a Keras model file (Keras3) or directory (Keras2) for the regression " + "of the primary particle arrival direction based on spherical coordinate offsets." + ), + allow_none=True, + exists=True, + directory_ok=True, + file_ok=True, + ).tag(config=True) + + batch_size = Int( + default_value=64, + allow_none=False, + help="Size of the batch to perform inference of the neural network.", + ).tag(config=True) + + output_path = Path( + default_value="./output.dl2.h5", + allow_none=False, + help="Output path to save the dl2 prediction results", + ).tag(config=True) + + overwrite_tables = Bool( + default_value=True, + allow_none=False, + help="Overwrite the table in the output file if it exists", + ).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", + ("e", "energy_model"): "PredictCTLearnModel.load_energy_model_from", + ( + "d", + "cameradirection_model", + ): "PredictCTLearnModel.load_cameradirection_model_from", + ("s", "skydirection_model"): "PredictCTLearnModel.load_skydirection_model_from", + ("o", "output"): "PredictCTLearnModel.output_path", + } + + flags = { + **flag( + "dl1-features", + "PredictCTLearnModel.dl1_features", + "Include dl1 features", + "Exclude dl1 features", + ), + **flag( + "dl2-telescope", + "PredictCTLearnModel.dl2_telescope", + "Include dl2 telescope-event-wise data in the output file", + "Exclude dl2 telescope-event-wise data in the output file", + ), + **flag( + "dl2-subarray", + "PredictCTLearnModel.dl2_subarray", + "Include dl2 telescope-event-wise data in the output file", + "Exclude dl2 telescope-event-wise data in the output file", + ), + **flag( + "use-HDF5Merger", + "PredictCTLearnModel.use_HDF5Merger", + "Copy data using the HDF5Merger component (CAUTION: This can not be used if the output file already exists)", + "Do not copy data using the HDF5Merger component", + ), + **flag( + "r0-waveforms", + "HDF5Merger.r0_waveforms", + "Include r0 waveforms", + "Exclude r0 waveforms", + ), + **flag( + "r1-waveforms", + "HDF5Merger.r1_waveforms", + "Include r1 waveforms", + "Exclude r1 waveforms", + ), + **flag( + "dl1-parameters", + "HDF5Merger.dl1_parameters", + "Include dl1 parameters", + "Exclude dl1 parameters", + ), + **flag( + "dl1-images", + "HDF5Merger.dl1_images", + "Include dl1 images", + "Exclude dl1 images", + ), + **flag( + "true-parameters", + "HDF5Merger.true_parameters", + "Include true parameters", + "Exclude true parameters", + ), + **flag( + "true-images", + "HDF5Merger.true_images", + "Include true images", + "Exclude true images", + ), + } + + classes = classes_with_traits(DLDataReader) + + def setup(self): + # Check if the ctapipe HDF5Merger component is enabled + if self.use_HDF5Merger: + if os.path.exists(self.output_path): + raise ToolConfigurationError( + f"The output file '{self.output_path}' already exists. Please use " + "'--no-use-HDF5Merger' to disable the usage of the HDF5Merger component." + ) + # Copy selected tables from the input file to the output file + self.log.info("Copying to output destination.") + with HDF5Merger(self.output_path, parent=self) as merger: + merger(self.input_url) + else: + self.log.info( + "No copy to output destination, since the usage of the HDF5Merger component is disabled." + ) + + if tf is None: + raise ImportError("TensorFlow is required for prediction. Install it with 'pip install ctlearn[tf]' or 'pip install ctlearn[all]'.") + + # 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) + + # Set up the data reader + self.log.info("Loading data reader:") + self.log.info("For a large dataset, this may take a while...") + self.dl1dh_reader = DLDataReader.from_name( + self.dl1dh_reader_type, + input_url_signal=[self.input_url], + parent=self, + ) + self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) + # Check if the number of events is enough to form a batch + if self.dl1dh_reader._get_n_events() < self.batch_size: + raise ToolConfigurationError( + f"{self.dl1dh_reader._get_n_events()} events are not enough " + f"to form a batch of size {self.batch_size}. Reduce the batch size." + ) + # 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 + ) + + def finish(self): + self.log.info("Tool is shutting down") + + def _predict_with_model(self, model_path): + """ + Load and predict with a CTLearn 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. + + Parameters + ---------- + model_path : str + Path to a Keras model file (Keras3) or directory (Keras2). + + Returns + ------- + predict_data : astropy.table.Table + Table containing the prediction results. + 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( + self.dl1dh_reader, + self.indices, + tasks=[], + batch_size=self.batch_size * self.strategy.num_replicas_in_sync, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + # Keras is only considering the last complete batch. + # In prediction mode we don't want to loose the last + # uncomplete batch, so we are creating an additional + # batch generator for the remaining events. + 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( + self.dl1dh_reader, + last_batch_indices, + tasks=[], + batch_size=self.last_batch_size, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + # Load the model from the specified path + model = keras.saving.load_model(model_path) + prediction_colname = ( + model.layers[-1].name if model.layers[-1].name != "softmax" else "type" + ) + backbone_model, feature_vectors = None, None + if self.dl1_features: + # Get the backbone model which is the second layer of the model + backbone_model = 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) + # Apply the backbone model with the data loader to retrieve the feature vectors + feature_vectors = backbone_model.predict( + data_loader, verbose=self.keras_verbose + ) + # Apply the head model with the feature vectors to retrieve the prediction + predict_data = Table( + { + prediction_colname: head.predict( + feature_vectors, verbose=self.keras_verbose + ) + } + ) + # 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 + ) + feature_vectors = np.concatenate( + (feature_vectors, feature_vectors_last_batch) + ) + predict_data = vstack( + [ + predict_data, + Table( + { + prediction_colname: head.predict( + feature_vectors_last_batch, + verbose=self.keras_verbose, + ) + } + ), + ] + ) + else: + # Predict the data using the loaded model + predict_data = model.predict(data_loader, verbose=self.keras_verbose) + # Create a astropy table with the prediction results + # The classification task has a softmax layer as the last layer + # which returns the probabilities for each class in an array, while + # the regression tasks have output neurons which returns the + # predicted value for the task in a dictionary. + if prediction_colname == "type": + predict_data = Table({prediction_colname: predict_data}) + else: + predict_data = Table(predict_data) + # 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 + ) + if model.layers[-1].name == "type": + predict_data_last_batch = Table( + {prediction_colname: predict_data_last_batch} + ) + else: + predict_data_last_batch = Table(predict_data_last_batch) + predict_data = vstack([predict_data, predict_data_last_batch]) + return predict_data, feature_vectors + + def _predict_classification(self, example_identifiers): + """ + Predict the classification of the primary particle type. + + This method uses a pre-trained type model to predict the type of the primary particle + for a given set of example identifiers. The predicted classification score ('gammaness') + is added to the example identifiers table. + + Parameters: + ----------- + classification_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + predicted classification score ('gammaness'). + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the classification of the primary particle type..." + ) + # Predict the data using the loaded type_model + predict_data, feature_vectors = self._predict_with_model( + self.load_type_model_from + ) + # Create prediction table and add the predicted classification score ('gammaness') + classification_table = example_identifiers.copy() + classification_table.add_column( + predict_data["type"].T[1], name=f"{self.prefix}_tel_prediction" + ) + return classification_table, feature_vectors + + def _predict_energy(self, example_identifiers): + """ + Predict the energy of the primary particle. + + This method uses a pre-trained energy model to predict the energy of the primary particle + for a given set of example identifiers. The predicted energy is then converted from + log10(TeV) to TeV and added to the example identifiers table. + + Parameters: + ----------- + energy_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed energy in TeV. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info("Predicting for the regression of the primary particle energy...") + # Predict the data using the loaded energy_model + predict_data, feature_vectors = self._predict_with_model( + self.load_energy_model_from + ) + # Convert the reconstructed energy from log10(TeV) to TeV + reco_energy = u.Quantity( + np.power(10, np.squeeze(predict_data["energy"])), + unit=u.TeV, + ) + # Create prediction table and add the reconstructed energy in TeV + energy_table = example_identifiers.copy() + energy_table.add_column(reco_energy, name=f"{self.prefix}_tel_energy") + return energy_table, feature_vectors + + def _predict_cameradirection(self, example_identifiers): + """ + Predict the arrival direction of the primary particle based on camera coordinate offsets. + + This method uses a pre-trained direction model to predict the arrival direction of the + primary particle for a given set of example identifiers. The predicted camera coordinate offsets + is added to the example identifiers table. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + + Returns: + -------- + cameradirection_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed camera coordinate offsets in x and y. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the regression of the primary particle arrival direction based on camera coordinate offsets..." + ) + # Predict the data using the loaded direction_model + predict_data, feature_vectors = self._predict_with_model( + self.load_cameradirection_model_from + ) + # For the direction task, the prediction is the camera coordinate offset in x and y + # from the telescope pointing. + cam_coord_offset_x = u.Quantity(predict_data["cameradirection"].T[0], unit=u.m) + cam_coord_offset_y = u.Quantity(predict_data["cameradirection"].T[1], unit=u.m) + # Create prediction table and add the reconstructed energy in TeV + cameradirection_table = example_identifiers.copy() + cameradirection_table.add_column(cam_coord_offset_x, name="cam_coord_offset_x") + cameradirection_table.add_column(cam_coord_offset_y, name="cam_coord_offset_y") + return cameradirection_table, feature_vectors + + def _predict_skydirection(self, example_identifiers): + """ + Predict the arrival direction of the primary particle based on spherical coordinate offsets. + + This method uses a pre-trained direction model to predict the arrival direction of the primary + particle for a given set of example identifiers. The predicted spherical coordinate offsets is + added to the example identifiers table. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + + Returns: + -------- + skydirection_table : astropy.table.Table + Table containing the example identifiers with an additional column for the + reconstructed spherical coordinate offsets in fov_lon and fov_lat. + feature_vectors : np.ndarray + Feature vectors extracted from the backbone model. + """ + self.log.info( + "Predicting for the regression of the primary particle arrival direction based on spherical coordinate offsets..." + ) + # Predict the data using the loaded direction_model + predict_data, feature_vectors = self._predict_with_model( + self.load_skydirection_model_from + ) + # For the direction task, the prediction is the spherical offset in fov_lon and fov_lat + # from the telescope pointing. + fov_lon = u.Quantity(predict_data["skydirection"].T[0], unit=u.deg) + fov_lat = u.Quantity(predict_data["skydirection"].T[1], unit=u.deg) + # Create prediction table and add the reconstructed energy in TeV + skydirection_table = example_identifiers.copy() + skydirection_table.add_column(fov_lon, name="fov_lon") + skydirection_table.add_column(fov_lat, name="fov_lat") + return skydirection_table, feature_vectors + + def _transform_cam_coord_offsets_to_sky(self, table) -> Table: + """ + Transform the predicted camera coordinate offsets w.r.t. the telescope pointing to Alt/Az coordinates. + + This method converts the predicted camera coordinate offsets w.r.t. the telescope pointing + in the provided table to Alt/Az coordinates. It also removes the unnecessary columns + from the table that do not the ctapipe DL2 data format. + + Parameters: + ----------- + table : astropy.table.Table + A Table containing the trigger time, telescope pointing, and predicted camera coordinate offsets. + + Returns: + -------- + table : astropy.table.Table + A Table with the Alt/Az coordinates following the ctapipe DL2 data format. + """ + # Get the telescope ID from the table + tel_id = table["tel_id"][0] + # Set the telescope position + tel_ground_frame = self.dl1dh_reader.subarray.tel_coords[ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] + # Set the trigger timestamp based on the process type + if self.dl1dh_reader.process_type == ProcessType.Simulation: + trigger_time = LST_EPOCH + elif self.dl1dh_reader.process_type == ProcessType.Observation: + trigger_time = table["time"] + # Set the telescope pointing with the trigger timestamp and the telescope position + altaz = AltAz( + location=tel_ground_frame.to_earth_location(), + obstime=trigger_time, + ) + # Set the telescope pointing + tel_pointing = SkyCoord( + az=table["pointing_azimuth"], + alt=table["pointing_altitude"], + frame=altaz, + ) + # Set the camera frame with the focal length and rotation of the camera + camera_frame = CameraFrame( + focal_length=self.dl1dh_reader.subarray.tel[ + tel_id + ].camera.geometry.frame.focal_length, + rotation=self.dl1dh_reader.pix_rotation[tel_id], + telescope_pointing=tel_pointing, + ) + # Set the camera coordinate offset + cam_coord_offset = SkyCoord( + x=table["cam_coord_offset_x"], + y=table["cam_coord_offset_y"], + frame=camera_frame, + ) + # tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] + # Transform the true Alt/Az coordinates to camera coordinates + reco_direction = cam_coord_offset.transform_to(altaz) + # Add the reconstructed direction (az, alt) to the prediction table + table.add_column(reco_direction.az.to(u.deg), name=f"{self.prefix}_tel_az") + table.add_column(reco_direction.alt.to(u.deg), name=f"{self.prefix}_tel_alt") + # Remove unnecessary columns from the table that do not the ctapipe DL2 data format + table.remove_columns( + [ + "time", + "pointing_azimuth", + "pointing_altitude", + "cam_coord_offset_x", + "cam_coord_offset_y", + ] + ) + return table + + def _transform_spher_coord_offsets_to_sky(self, table) -> Table: + """ + Transform the predicted spherical offsets w.r.t. the telescope pointing to Alt/Az coordinates. + + This method converts the predicted spherical offsets w.r.t. the telescope pointing + in the provided table to Alt/Az coordinates. It also removes the unnecessary columns + from the table that do not the ctapipe DL2 data format. + + Parameters: + ----------- + table : astropy.table.Table + A Table containing the trigger time, telescope pointing, and predicted spherical offsets. + + Returns: + -------- + table : astropy.table.Table + A Table with the Alt/Az coordinates following the ctapipe DL2 data format. + """ + + # Set the trigger timestamp based on the process type + if self.dl1dh_reader.process_type == ProcessType.Simulation: + trigger_time = LST_EPOCH + elif self.dl1dh_reader.process_type == ProcessType.Observation: + trigger_time = table["time"] + # Set the AltAz frame with the reference location and time + altaz = AltAz( + location=self.dl1dh_reader.subarray.reference_location, + obstime=trigger_time, + ) + # Set the array pointing + array_pointing = SkyCoord( + az=table["pointing_azimuth"], + alt=table["pointing_altitude"], + frame=altaz, + ) + # Set the nominal frame with the array pointing + nom_frame = NominalFrame( + origin=array_pointing, + location=self.dl1dh_reader.subarray.reference_location, + obstime=trigger_time, + ) + # Set the reco direction in (fov_lon, fov_lat) coordinates + reco_direction = SkyCoord( + fov_lon=table["fov_lon"], + fov_lat=table["fov_lat"], + frame=nom_frame, + ) + # Transform the reco direction from nominal frame to the AltAz frame + sky_coord = reco_direction.transform_to(altaz) + # Add the reconstructed direction (az, alt) to the prediction table + table.add_column(sky_coord.az.to(u.deg), name=f"{self.prefix}_az") + table.add_column(sky_coord.alt.to(u.deg), name=f"{self.prefix}_alt") + # Remove unnecessary columns from the table that do not the ctapipe DL2 data format + table.remove_columns( + [ + "time", + "pointing_azimuth", + "pointing_altitude", + "fov_lon", + "fov_lat", + ] + ) + return table + + def _create_nan_table(self, nonexample_identifiers, columns, shapes): + """ + Create a table with NaNs for missing predictions. + + This method creates a table with NaNs for missing predictions for the non-example identifiers. + In stereo mode, the table also a column for the valid telescopes is added with all False values. + + Parameters: + ----------- + nonexample_identifiers : astropy.table.Table + Table containing the non-example identifiers. + columns : list of str + List of column names to create in the table. + shapes : list of shapes + List of shapes for the columns to create in the table. + + Returns: + -------- + nan_table : astropy.table.Table + Table containing NaNs for missing predictions. + """ + # Create a table with NaNs for missing predictions + nan_table = nonexample_identifiers.copy() + for column_name, shape in zip(columns, shapes): + nan_table.add_column(np.full(shape, np.nan), name=column_name) + # Add that no telescope is valid for the non-example identifiers in stereo mode + if self.dl1dh_reader.mode == "stereo": + nan_table.add_column( + np.zeros( + (len(nonexample_identifiers), len(self.dl1dh_reader.tel_ids)), + dtype=bool, + ), + name=f"{self.prefix}_telescopes", + ) + return nan_table + + def _store_pointing(self, all_identifiers): + """ + Store the telescope pointing table from to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the telescope pointing information. + """ + + # Initialize the pointing interpolator from ctapipe + pointing_interpolator = PointingInterpolator( + bounds_error=False, extrapolate=True + ) + pointing_info = [] + for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: + # Get the telescope pointing from the dl1dh reader + tel_pointing = self.dl1dh_reader.telescope_pointings[f"tel_{tel_id:03d}"] + # Add the telescope pointing table to the pointing interpolator + pointing_interpolator.add_table(tel_id, tel_pointing) + tel_identifiers = all_identifiers.copy() + if self.dl1dh_reader.mode == "mono": + tel_identifiers = tel_identifiers[tel_identifiers["tel_id"] == tel_id] + # Interpolate the telescope pointing + tel_altitude, tel_azimuth = pointing_interpolator( + tel_id, tel_identifiers["time"] + ) + tel_identifiers.add_column(tel_azimuth, name="pointing_azimuth") + tel_identifiers.add_column(tel_altitude, name="pointing_altitude") + pointing_info.append(tel_identifiers) + if self.dl1dh_reader.mode == "mono": + tel_pointing_table = Table( + { + "time": tel_identifiers["time"], + "azimuth": tel_identifiers["pointing_azimuth"], + "altitude": tel_identifiers["pointing_altitude"], + } + ) + write_table( + tel_pointing_table, + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 telescope pointing table was stored in '%s' under '%s'", + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + ) + pointing_info = vstack(pointing_info) + if self.dl1dh_reader.mode == "stereo": + # Group the pointing information by subarray event keys + # TODO: This needs to be debugged with SST1M data + pointing_info_grouped = pointing_info.group_by(SUBARRAY_EVENT_KEYS) + pointing_mean = pointing_info_grouped.groups.aggregate(np.mean) + pointing_info = join( + all_identifiers, + pointing_mean, + keys=SUBARRAY_EVENT_KEYS, + ) + # TODO: use keep_order for astropy v7.0.0 + pointing_info.sort(SUBARRAY_EVENT_KEYS) + # Create the pointing table + pointing_table = Table( + { + "time": pointing_info["time"], + "array_azimuth": pointing_info["pointing_azimuth"], + "array_altitude": pointing_info["pointing_altitude"], + "array_ra": np.nan * np.ones(len(pointing_info)), + "array_dec": np.nan * np.ones(len(pointing_info)), + } + ) + # Save the pointing table to the output file + write_table( + pointing_table, + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 subarray pointing table was stored in '%s' under '%s'", + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + ) + return pointing_info + + def _create_feature_vectors_table( + self, + example_identifiers, + nonexample_identifiers=None, + classification_feature_vectors=None, + energy_feature_vectors=None, + direction_feature_vectors=None, + ): + """ + Create the table for the DL1 feature vectors. + + This method creates a table with the DL1 feature vectors for the example identifiers and fill NaNs for + non-example identifiers. The feature vectors are stored in the columns of the table. The table also + contains a column for the valid predictions. + + Parameters: + ----------- + example_identifiers : astropy.table.Table + Table containing the example identifiers. + nonexample_identifiers : astropy.table.Table or None + Table containing the non-example identifiers to fill the NaNs. + classification_feature_vectors : np.ndarray or None + Array containing the classification feature vectors. + energy_feature_vectors : np.ndarray or None + Array containing the energy feature vectors. + direction_feature_vectors : np.ndarray or None + Array containing the direction feature vectors. + + Returns: + -------- + feature_vector_table : astropy.table.Table + Table containing the DL1 feature vectors for the example and non-example identifiers. + """ + # Create the feature vector table + feature_vector_table = example_identifiers.copy() + feature_vector_table.remove_columns( + ["pointing_azimuth", "pointing_altitude", "time"] + ) + columns_list, shapes_list = [], [] + if classification_feature_vectors is not None: + is_valid_col = ~np.isnan( + np.min(classification_feature_vectors, axis=1), dtype=bool + ) + feature_vector_table.add_column( + classification_feature_vectors, + name=f"{self.prefix}_tel_classification_feature_vectors", + ) + if nonexample_identifiers is not None: + columns_list.append(f"{self.prefix}_tel_classification_feature_vectors") + shapes_list.append( + ( + len(nonexample_identifiers), + classification_feature_vectors.shape[1], + ) + ) + if energy_feature_vectors is not None: + is_valid_col = ~np.isnan(np.min(energy_feature_vectors, axis=1), dtype=bool) + feature_vector_table.add_column( + energy_feature_vectors, name=f"{self.prefix}_tel_energy_feature_vectors" + ) + if nonexample_identifiers is not None: + columns_list.append(f"{self.prefix}_tel_energy_feature_vectors") + shapes_list.append( + ( + len(nonexample_identifiers), + energy_feature_vectors.shape[1], + ) + ) + if direction_feature_vectors is not None: + is_valid_col = ~np.isnan( + np.min(direction_feature_vectors, axis=1), dtype=bool + ) + feature_vector_table.add_column( + direction_feature_vectors, + name=f"{self.prefix}_tel_geometry_feature_vectors", + ) + if nonexample_identifiers is not None: + columns_list.append(f"{self.prefix}_tel_geometry_feature_vectors") + shapes_list.append( + ( + len(nonexample_identifiers), + direction_feature_vectors.shape[1], + ) + ) + # Produce output table with NaNs for missing predictions + if nonexample_identifiers is not None: + if len(nonexample_identifiers) > 0: + nan_table = self._create_nan_table( + nonexample_identifiers, + columns=columns_list, + shapes=shapes_list, + ) + feature_vector_table = vstack([feature_vector_table, nan_table]) + is_valid_col = np.concatenate( + (is_valid_col, np.zeros(len(nonexample_identifiers), dtype=bool)) + ) + # Add is_valid column to the feature vector table + feature_vector_table.add_column( + is_valid_col, + name=f"{self.prefix}_tel_is_valid", + ) + return feature_vector_table + + +class MonoPredictCTLearnModel(PredictCTLearnModel): + """ + Tool to predict the gammaness, energy and arrival direction from monoscopic R1/DL1 data using CTLearn models. + + This tool extends the ``PredictCTLearnModel`` to specifically handle monoscopic R1/DL1 data. The prediction + is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. + It also stores the telescope pointing monitoring and DL1 feature vectors (if selected) in the output file. + + Attributes + ---------- + name : str + Name of the tool. + description : str + Description of the tool. + examples : str + Examples of how to use the tool. + + Methods + ------- + start() + Start the tool. + _store_mc_telescope_pointing(all_identifiers) + Store the telescope pointing table for the mono mode for MC simulation. + """ + + name = "ctlearn-predict-mono-model" + description = __doc__ + + examples = """ + To predict from pixel-wise image data in mono mode using trained CTLearn models: + > ctlearn-predict-mono-model \\ + --input_url input.dl1.h5 \\ + --PredictCTLearnModel.batch_size=64 \\ + --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ + --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" \\ + --dl1-features \\ + --use-HDF5Merger \\ + --no-dl1-images \\ + --no-true-images \\ + --output output.dl2.h5 \\ + --PredictCTLearnModel.overwrite_tables=True \\ + + To predict from pixel-wise waveform data in mono mode using trained CTLearn models: + > ctlearn-predict-mono-model \\ + --input_url input.r1.h5 \\ + --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" \\ + --use-HDF5Merger \\ + --no-r0-waveforms \\ + --no-r1-waveforms \\ + --no-dl1-images \\ + --no-true-images \\ + --output output.dl2.h5 \\ + --PredictCTLearnModel.overwrite_tables=True \\ + """ + + stereo_combiner_cls = ComponentName( + StereoCombiner, + default_value="StereoMeanCombiner", + help="Which stereo combination method to use after the monoscopic reconstruction.", + ).tag(config=True) + + def start(self): + self.log.info("Processing the telescope pointings...") + # Retrieve the IDs from the dl1dh for the prediction tables + example_identifiers = self.dl1dh_reader.example_identifiers.copy() + example_identifiers.keep_columns(TELESCOPE_EVENT_KEYS) + all_identifiers = self.dl1dh_reader.tel_trigger_table.copy() + all_identifiers.keep_columns(TELESCOPE_EVENT_KEYS + ["time"]) + nonexample_identifiers = setdiff( + all_identifiers, example_identifiers, keys=TELESCOPE_EVENT_KEYS + ) + nonexample_identifiers.remove_column("time") + # Pointing table for the mono mode for MC simulation + if self.dl1dh_reader.process_type == ProcessType.Simulation: + pointing_info = self._store_mc_telescope_pointing(all_identifiers) + + # Pointing table for the observation mode + if self.dl1dh_reader.process_type == ProcessType.Observation: + pointing_info = super()._store_pointing(all_identifiers) + + self.log.info("Starting the prediction...") + classification_feature_vectors = None + if self.load_type_model_from is not None: + self.type_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefix, + property=ReconstructionProperty.PARTICLE_TYPE, + parent=self, + ) + # Predict the energy of the primary particle + classification_table, classification_feature_vectors = ( + super()._predict_classification(example_identifiers) + ) + if self.dl2_telescope: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_prediction"], + shapes=[(len(nonexample_identifiers),)], + ) + classification_table = vstack([classification_table, nan_table]) + # Add is_valid column to the energy table + classification_table.add_column( + ~np.isnan( + classification_table[f"{self.prefix}_tel_prediction"].data, + dtype=bool, + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + classification_table, + ParticleClassificationContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = classification_table["tel_id"] == tel_id + classification_tel_table = classification_table[telescope_mask] + classification_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Save the prediction to the output file for the selected telescope + write_table( + classification_tel_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/classification/{self.prefix}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info("Processing and storing the subarray type prediction...") + # Combine the telescope predictions to the subarray prediction using the stereo combiner + subarray_classification_table = self.type_stereo_combiner.predict_table( + classification_table + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if ( + subarray_classification_table[f"{self.prefix}_telescopes"].dtype + != np.bool_ + ): + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + ( + len(subarray_classification_table), + len(self.dl1dh_reader.tel_ids), + ), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + subarray_classification_table[f"{self.prefix}_telescopes"] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] = True + # Overwrite the column with the boolean mask with fix length + subarray_classification_table[f"{self.prefix}_telescopes"] = ( + reco_telescopes + ) + # Save the prediction to the output file + write_table( + subarray_classification_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + ) + energy_feature_vectors = None + if self.load_energy_model_from is not None: + self.energy_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefix, + property=ReconstructionProperty.ENERGY, + parent=self, + ) + # Predict the energy of the primary particle + energy_table, energy_feature_vectors = super()._predict_energy( + example_identifiers + ) + if self.dl2_telescope: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_energy"], + shapes=[(len(nonexample_identifiers),)], + ) + energy_table = vstack([energy_table, nan_table]) + # Add is_valid column to the energy table + energy_table.add_column( + ~np.isnan( + energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + energy_table, + ReconstructedEnergyContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = energy_table["tel_id"] == tel_id + energy_tel_table = energy_table[telescope_mask] + energy_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Save the prediction to the output file + write_table( + energy_tel_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/energy/{self.prefix}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info( + "Processing and storing the subarray energy prediction..." + ) + # Combine the telescope predictions to the subarray prediction using the stereo combiner + subarray_energy_table = self.energy_stereo_combiner.predict_table( + energy_table + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if subarray_energy_table[f"{self.prefix}_telescopes"].dtype != np.bool_: + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + (len(subarray_energy_table), len(self.dl1dh_reader.tel_ids)), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + subarray_energy_table[f"{self.prefix}_telescopes"] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] = True + # Overwrite the column with the boolean mask with fix length + subarray_energy_table[f"{self.prefix}_telescopes"] = reco_telescopes + # Save the prediction to the output file + write_table( + subarray_energy_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + ) + direction_feature_vectors = None + if self.load_cameradirection_model_from is not None: + self.geometry_stereo_combiner = StereoCombiner.from_name( + self.stereo_combiner_cls, + prefix=self.prefix, + property=ReconstructionProperty.GEOMETRY, + parent=self, + ) + # Join the prediction table with the telescope pointing table + example_identifiers = join( + left=example_identifiers, + right=pointing_info, + keys=TELESCOPE_EVENT_KEYS, + ) + # Predict the arrival direction of the primary particle + direction_table, direction_feature_vectors = ( + super()._predict_cameradirection(example_identifiers) + ) + direction_tel_tables = [] + if self.dl2_telescope: + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = direction_table["tel_id"] == tel_id + direction_tel_table = direction_table[telescope_mask] + direction_tel_table = super()._transform_cam_coord_offsets_to_sky( + direction_tel_table + ) + # Produce output table with NaNs for missing predictions + nan_telescope_mask = nonexample_identifiers["tel_id"] == tel_id + nonexample_identifiers_tel = nonexample_identifiers[ + nan_telescope_mask + ] + if len(nonexample_identifiers_tel) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers_tel, + columns=[f"{self.prefix}_tel_alt", f"{self.prefix}_tel_az"], + shapes=[ + (len(nonexample_identifiers_tel),), + (len(nonexample_identifiers_tel),), + ], + ) + direction_tel_table = vstack([direction_tel_table, nan_table]) + direction_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Add is_valid column to the direction table + direction_tel_table.add_column( + ~np.isnan( + direction_tel_table[f"{self.prefix}_tel_alt"].data, + dtype=bool, + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Add the default values and meta data to the table + add_defaults_and_meta( + direction_tel_table, + ReconstructedGeometryContainer, + prefix=self.prefix, + add_tel_prefix=True, + ) + direction_tel_tables.append(direction_tel_table) + # Save the prediction to the output file + write_table( + direction_tel_table, + self.output_path, + f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_TELESCOPE_GROUP}/geometry/{self.prefix}/tel_{tel_id:03d}", + ) + if self.dl2_subarray: + self.log.info( + "Processing and storing the subarray geometry prediction..." + ) + # Stack the telescope tables to the subarray table + direction_tel_tables = vstack(direction_tel_tables) + # Sort the table by the telescope event keys + direction_tel_tables.sort(TELESCOPE_EVENT_KEYS) + # Combine the telescope predictions to the subarray prediction using the stereo combiner + subarray_direction_table = self.geometry_stereo_combiner.predict_table( + direction_tel_tables + ) + # TODO: Remove temporary fix once the stereo combiner returns correct table + # Check if the table has to be converted to a boolean mask + if ( + subarray_direction_table[f"{self.prefix}_telescopes"].dtype + != np.bool_ + ): + # Create boolean mask for telescopes that participate in the stereo reconstruction combination + reco_telescopes = np.zeros( + (len(subarray_direction_table), len(self.dl1dh_reader.tel_ids)), + dtype=bool, + ) + # Loop over the table and set the boolean mask for the telescopes + for index, tel_id_mask in enumerate( + subarray_direction_table[f"{self.prefix}_telescopes"] + ): + if not tel_id_mask: + continue + for tel_id in tel_id_mask: + reco_telescopes[index][ + self.dl1dh_reader.subarray.tel_ids_to_indices(tel_id) + ] = True + # Overwrite the column with the boolean mask with fix length + subarray_direction_table[f"{self.prefix}_telescopes"] = ( + reco_telescopes + ) + # Save the prediction to the output file + write_table( + subarray_direction_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + ) + # Create the feature vector table if the DL1 features are enabled + if self.dl1_features: + self.log.info("Processing and storing dl1 feature vectors...") + feature_vector_table = super()._create_feature_vectors_table( + example_identifiers, + nonexample_identifiers, + classification_feature_vectors, + energy_feature_vectors, + direction_feature_vectors, + ) + # Loop over the selected telescopes and store the feature vectors + # for each telescope in the output file. The feature vectors are stored + # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. + for tel_id in self.dl1dh_reader.selected_telescopes[ + self.dl1dh_reader.tel_type + ]: + # Retrieve the example identifiers for the selected telescope + telescope_mask = feature_vector_table["tel_id"] == tel_id + feature_vectors_tel_table = feature_vector_table[telescope_mask] + feature_vectors_tel_table.sort(TELESCOPE_EVENT_KEYS) + # Save the prediction to the output file + write_table( + feature_vectors_tel_table, + self.output_path, + f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 feature vectors was stored in '%s' under '%s'", + self.output_path, + f"{DL1_TELESCOPE_GROUP}/features/{self.prefix}/tel_{tel_id:03d}", + ) + + def _store_mc_telescope_pointing(self, all_identifiers): + """ + Store the telescope pointing table from MC simulation to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the telescope pointing information. + """ + # Create the pointing table for each telescope + pointing_info = [] + for tel_id in self.dl1dh_reader.selected_telescopes[self.dl1dh_reader.tel_type]: + # Pointing table for the mono mode + tel_pointing = self.dl1dh_reader.get_tel_pointing(self.input_url, tel_id) + tel_pointing.rename_column("telescope_pointing_azimuth", "pointing_azimuth") + tel_pointing.rename_column( + "telescope_pointing_altitude", "pointing_altitude" + ) + # Join the prediction table with the telescope pointing table + tel_pointing = join( + left=tel_pointing, + right=all_identifiers, + keys=["obs_id", "tel_id"], + ) + # TODO: use keep_order for astropy v7.0.0 + tel_pointing.sort(TELESCOPE_EVENT_KEYS) + # Retrieve the example identifiers for the selected telescope + tel_pointing_table = Table( + { + "time": tel_pointing["time"], + "azimuth": tel_pointing["pointing_azimuth"], + "altitude": tel_pointing["pointing_altitude"], + } + ) + write_table( + tel_pointing_table, + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 telescope pointing table was stored in '%s' under '%s'", + self.output_path, + f"{POINTING_GROUP}/tel_{tel_id:03d}", + ) + pointing_info.append(tel_pointing) + pointing_info = vstack(pointing_info) + return pointing_info + + +class StereoPredictCTLearnModel(PredictCTLearnModel): + """ + Tool to predict the gammaness, energy and arrival direction from R1/DL1 stereoscopic data using CTLearn models. + + This tool extends the ``PredictCTLearnModel`` to specifically handle stereoscopic R1/DL1 data. The prediction + is performed using the CTLearn models. The data is stored in the output file following the ctapipe DL2 data format. + It also stores the telescope/subarray pointing monitoring and DL1 feature vectors (if selected) in the output file. + + Attributes + ---------- + name : str + Name of the tool. + description : str + Description of the tool. + examples : str + Examples of how to use the tool. + + Methods + ------- + start() + Start the tool. + _store_mc_subarray_pointing(all_identifiers) + Store the subarray pointing table for the stereo mode for MC simulation. + """ + + name = "ctlearn-predict-stereo-model" + description = __doc__ + + examples = """ + To predict from pixel-wise image data in stereo mode using trained CTLearn models: + > ctlearn-predict-stereo-model \\ + --input_url input.dl1.h5 \\ + --PredictCTLearnModel.batch_size=16 \\ + --PredictCTLearnModel.dl1dh_reader_type=DLImageReader \\ + --DLImageReader.channels=cleaned_image \\ + --DLImageReader.channels=cleaned_relative_peak_time \\ + --DLImageReader.image_mapper_type=BilinearMapper \\ + --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" \\ + --output output.dl2.h5 \\ + --PredictCTLearnModel.overwrite_tables=True \\ + """ + + def start(self): + self.log.info("Processing the telescope pointings...") + # Retrieve the IDs from the dl1dh for the prediction tables + example_identifiers = self.dl1dh_reader.unique_example_identifiers.copy() + example_identifiers.keep_columns(SUBARRAY_EVENT_KEYS) + all_identifiers = self.dl1dh_reader.subarray_trigger_table.copy() + all_identifiers.keep_columns(SUBARRAY_EVENT_KEYS + ["time"]) + nonexample_identifiers = setdiff( + all_identifiers, example_identifiers, keys=SUBARRAY_EVENT_KEYS + ) + nonexample_identifiers.remove_column("time") + # Construct the survival telescopes for each event of the example_identifiers + survival_telescopes = [] + for subarray_event in self.dl1dh_reader.example_identifiers_grouped.groups: + survival_mask = np.zeros(len(self.dl1dh_reader.tel_ids), dtype=bool) + survival_tels = [ + self.dl1dh_reader.subarray.tel_indices[tel_id] + for tel_id in subarray_event["tel_id"].data + ] + survival_mask[survival_tels] = True + survival_telescopes.append(survival_mask) + # Add the survival telescopes to the example_identifiers + example_identifiers.add_column( + survival_telescopes, name=f"{self.prefix}_telescopes" + ) + # Pointing table for the stereo mode for MC simulation + if self.dl1dh_reader.process_type == ProcessType.Simulation: + pointing_info = self._store_mc_subarray_pointing(all_identifiers) + + # Pointing table for the observation mode + if self.dl1dh_reader.process_type == ProcessType.Observation: + pointing_info = super()._store_pointing(all_identifiers) + + self.log.info("Starting the prediction...") + classification_feature_vectors = None + if self.load_type_model_from is not None: + # Predict the energy of the primary particle + classification_table, classification_feature_vectors = ( + super()._predict_classification(example_identifiers) + ) + if self.dl2_subarray: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_prediction"], + shapes=[(len(nonexample_identifiers),)], + ) + classification_table = vstack([classification_table, nan_table]) + # Add is_valid column to the energy table + classification_table.add_column( + ~np.isnan( + classification_table[f"{self.prefix}_tel_prediction"].data, + dtype=bool, + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Rename the columns for the stereo mode + classification_table.rename_column( + f"{self.prefix}_tel_prediction", f"{self.prefix}_prediction" + ) + classification_table.rename_column( + f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" + ) + classification_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + classification_table, + ParticleClassificationContainer, + prefix=self.prefix, + ) + # Save the prediction to the output file + write_table( + classification_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/classification/{self.prefix}", + ) + energy_feature_vectors = None + if self.load_energy_model_from is not None: + # Predict the energy of the primary particle + energy_table, energy_feature_vectors = super()._predict_energy( + example_identifiers + ) + if self.dl2_subarray: + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_tel_energy"], + shapes=[(len(nonexample_identifiers),)], + ) + energy_table = vstack([energy_table, nan_table]) + # Add is_valid column to the energy table + energy_table.add_column( + ~np.isnan( + energy_table[f"{self.prefix}_tel_energy"].data, dtype=bool + ), + name=f"{self.prefix}_tel_is_valid", + ) + # Rename the columns for the stereo mode + energy_table.rename_column( + f"{self.prefix}_tel_energy", f"{self.prefix}_energy" + ) + energy_table.rename_column( + f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" + ) + energy_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + energy_table, + ReconstructedEnergyContainer, + prefix=self.prefix, + ) + # Save the prediction to the output file + write_table( + energy_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/energy/{self.prefix}", + ) + direction_feature_vectors = None + if self.load_skydirection_model_from is not None: + # Join the prediction table with the telescope pointing table + example_identifiers = join( + left=example_identifiers, + right=pointing_info, + keys=SUBARRAY_EVENT_KEYS, + ) + # Predict the arrival direction of the primary particle + direction_table, direction_feature_vectors = super()._predict_skydirection( + example_identifiers + ) + if self.dl2_subarray: + # Transform the spherical coordinate offsets to sky coordinates + direction_table = super()._transform_spher_coord_offsets_to_sky( + direction_table + ) + # Produce output table with NaNs for missing predictions + if len(nonexample_identifiers) > 0: + nan_table = super()._create_nan_table( + nonexample_identifiers, + columns=[f"{self.prefix}_alt", f"{self.prefix}_az"], + shapes=[ + (len(nonexample_identifiers),), + (len(nonexample_identifiers),), + ], + ) + direction_table = vstack([direction_table, nan_table]) + # Add is_valid column to the direction table + direction_table.add_column( + ~np.isnan(direction_table[f"{self.prefix}_alt"].data, dtype=bool), + name=f"{self.prefix}_is_valid", + ) + direction_table.sort(SUBARRAY_EVENT_KEYS) + # Add the default values and meta data to the table + add_defaults_and_meta( + direction_table, + ReconstructedGeometryContainer, + prefix=self.prefix, + ) + # Save the prediction to the output file + write_table( + direction_table, + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL2 prediction data was stored in '%s' under '%s'", + self.output_path, + f"{DL2_SUBARRAY_GROUP}/geometry/{self.prefix}", + ) + + # Create the feature vector table if the DL1 features are enabled + if self.dl1_features: + self.log.info("Processing and storing dl1 feature vectors...") + feature_vector_table = super()._create_feature_vectors_table( + example_identifiers, + nonexample_identifiers, + classification_feature_vectors, + energy_feature_vectors, + direction_feature_vectors, + ) + # Loop over the selected telescopes and store the feature vectors + # for each telescope in the output file. The feature vectors are stored + # in the DL1_TELESCOPE_GROUP/features/{prefix}/tel_{tel_id:03d} table. + # Rename the columns for the stereo mode + feature_vector_table.rename_column( + f"{self.prefix}_tel_classification_feature_vectors", + f"{self.prefix}_classification_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefix}_tel_energy_feature_vectors", + f"{self.prefix}_energy_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefix}_tel_geometry_feature_vectors", + f"{self.prefix}_geometry_feature_vectors", + ) + feature_vector_table.rename_column( + f"{self.prefix}_tel_is_valid", f"{self.prefix}_is_valid" + ) + feature_vector_table.sort(SUBARRAY_EVENT_KEYS) + # Save the prediction to the output file + write_table( + feature_vector_table, + self.output_path, + f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 feature vectors was stored in '%s' under '%s'", + self.output_path, + f"{DL1_SUBARRAY_GROUP}/features/{self.prefix}", + ) + + def _store_mc_subarray_pointing(self, all_identifiers): + """ + Store the subarray pointing table from MC simulation to the output file. + + Parameters: + ----------- + all_identifiers : astropy.table.Table + Table containing the subarray pointing information. + """ + # Read the subarray pointing table + pointing_info = read_table( + self.input_url, + f"{SIMULATION_CONFIG_TABLE}", + ) + # Assuming min_az = max_az and min_alt = max_alt + pointing_info.keep_columns(["obs_id", "min_az", "min_alt"]) + pointing_info.rename_column("min_az", "pointing_azimuth") + pointing_info.rename_column("min_alt", "pointing_altitude") + # Join the prediction table with the telescope pointing table + pointing_info = join( + left=pointing_info, + right=all_identifiers, + keys=["obs_id"], + ) + # TODO: use keep_order for astropy v7.0.0 + pointing_info.sort(SUBARRAY_EVENT_KEYS) + # Create the pointing table + pointing_table = Table( + { + "time": pointing_info["time"], + "array_azimuth": pointing_info["pointing_azimuth"], + "array_altitude": pointing_info["pointing_altitude"], + "array_ra": np.nan * np.ones(len(pointing_info)), + "array_dec": np.nan * np.ones(len(pointing_info)), + } + ) + # Save the pointing table to the output file + write_table( + pointing_table, + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + overwrite=self.overwrite_tables, + ) + self.log.info( + "DL1 subarray pointing table was stored in '%s' under '%s'", + self.output_path, + f"{SUBARRAY_POINTING_GROUP}", + ) + return pointing_info + + +def mono_tool(): + # Run the tool + mono_tool = MonoPredictCTLearnModel() + mono_tool.run() + + +def stereo_tool(): + # Run the tool + stereo_tool = StereoPredictCTLearnModel() + stereo_tool.run() + + +if __name__ == "mono_tool": + mono_tool() + +if __name__ == "stereo_tool": + stereo_tool() \ 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..da72dded 100644 --- a/ctlearn/tools/tests/test_predict_LST1.py +++ b/ctlearn/tools/tests/test_predict_LST1.py @@ -2,6 +2,8 @@ import numpy as np import pytest +pytest.importorskip("tensorflow") + from ctapipe.core import run_tool from ctapipe.io import TableLoader from ctlearn.tools import LST1PredictionTool diff --git a/ctlearn/tools/tests/test_predict_model.py b/ctlearn/tools/tests/test_predict_model.py index 5f17fc34..81093c6a 100644 --- a/ctlearn/tools/tests/test_predict_model.py +++ b/ctlearn/tools/tests/test_predict_model.py @@ -2,6 +2,8 @@ import numpy as np import pytest +pytest.importorskip("tensorflow") + from ctapipe.core import run_tool from ctapipe.io import TableLoader from ctlearn.tools import MonoPredictCTLearnModel, StereoPredictCTLearnModel diff --git a/ctlearn/tools/tests/test_train_model.py b/ctlearn/tools/tests/test_train_model.py index d2291652..71f3f6fb 100644 --- a/ctlearn/tools/tests/test_train_model.py +++ b/ctlearn/tools/tests/test_train_model.py @@ -2,12 +2,16 @@ import pytest import shutil +pytest.importorskip("tensorflow") +from unittest import mock + from ctapipe.core import run_tool -from ctlearn.tools import TrainCTLearnModel +from ctlearn.tools import DLFrameWork @pytest.mark.parametrize("reco_task", ["type", "energy", "cameradirection"]) -def test_train_ctlearn_model(reco_task, dl1_gamma_file, dl1_proton_file, tmp_path): +@mock.patch("ctapipe.instrument.SubarrayDescription.__eq__", return_value=True) +def test_train_ctlearn_model(mock_eq, reco_task, dl1_gamma_file, dl1_proton_file, tmp_path): """ Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks. Each test run gets its own isolated temp directories. @@ -40,18 +44,16 @@ def test_train_ctlearn_model(reco_task, dl1_gamma_file, dl1_proton_file, tmp_pat f"--DLImageReader.allowed_tels={allowed_tels}", ] - # Include background only for classification task if reco_task == "type": argv.extend( [ f"--background={background_dir}", "--pattern-background=*.dl1.h5", - "--DLImageReader.enforce_subarray_equality=False", ] ) # Run training - assert run_tool(TrainCTLearnModel(), argv=argv, cwd=tmp_path) == 0 + assert run_tool(DLFrameWork(), argv=argv, cwd=tmp_path) == 0 # --- Additional checks --- # Check that the trained model exists diff --git a/ctlearn/tools/train/_version.py b/ctlearn/tools/train/_version.py new file mode 100644 index 00000000..b8023d8b --- /dev/null +++ b/ctlearn/tools/train/_version.py @@ -0,0 +1 @@ +__version__ = '0.0.1' diff --git a/ctlearn/tools/train/base_train_model.py b/ctlearn/tools/train/base_train_model.py new file mode 100644 index 00000000..a02b34d0 --- /dev/null +++ b/ctlearn/tools/train/base_train_model.py @@ -0,0 +1,603 @@ +import numpy as np +import shutil +from ctapipe.core import Tool +from ctapipe.core.tool import ToolConfigurationError +from ctapipe.core.traits import ( + Bool, + CaselessStrEnum, + Path, + Float, + Int, + List, + Dict, + classes_with_traits, + ComponentName, + Unicode, +) +from dl1_data_handler.reader import DLDataReader, TableQualityQuery + +class TrainCTLearnModel(Tool): + """ + Base class for training a ``CTLearnModel`` on R1/DL1a data using the ``DLDataReader`` and ``DLDataLoader``. + + 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 ``DLDataReader`` and ``DLDataLoader``. The ``start`` method + is implemented in the subclasses to train the model using the specified framework (Keras or PyTorch). + 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 + + Attributes + ---------- + input_dir_signal : Path + Input directory for signal events. + file_pattern_signal : List[Unicode] + List of specific file pattern for matching files in ``input_dir_signal``. + input_dir_background : Path + Input directory for background events. + file_pattern_background : List[Unicode] + List of specific file pattern for matching files in ``input_dir_background``. + dl1dh_reader_type : ComponentName + Type of the DLDataReader to use for reading the input data. + dl1dh_reader : DLDataReader + DLDataReader instance for reading the input data. + stack_telescope_images : Bool + Set whether to stack the telescope images in the data loader. + sort_by_intensity : Bool + Set whether to sort the telescope images by intensity in the data loader. + output_dir : Path + Output directory for the trained reconstructor. + reco_tasks : List[Unicode] + List of reconstruction tasks to perform. + n_epochs : Int + Number of epochs to train the neural network. + batch_size : Int + Size of the batch to train the neural network. + validation_split : Float + Fraction of the data to use for validation. + optimizer : Dict + Optimizer to use for training. + random_seed : Int + Random seed for shuffling the data before the training/validation split and after the end of an epoch. + save_onnx : Bool + Set whether to save model in an ONNX file. + overwrite : Bool + Overwrite output dir if it exists. + + Methods + ------- + setup() + Set up the data reader and data loaders for training and validation. + finish() + Save the trained model in the output directory in ONNX if selected. + """ + + name = "ctlearn-train-model-base" + + framework_type = CaselessStrEnum( + ["pytorch", "keras"], + default_value="keras", + help="Framework to use pytorch or keras", + ).tag(config=True) + + input_dir_signal = Path( + help="Input directory for signal events", + allow_none=False, + exists=True, + directory_ok=True, + file_ok=False, + ).tag(config=True) + + file_pattern_signal = List( + trait=Unicode(), + default_value=["*.h5"], + help="List of specific file pattern for matching files in ``input_dir_signal``", + ).tag(config=True) + + input_dir_background = Path( + default_value=None, + help="Input directory for background events", + allow_none=True, + exists=True, + directory_ok=True, + file_ok=False, + ).tag(config=True) + + file_pattern_background = List( + trait=Unicode(), + default_value=["*.h5"], + help="List of specific file pattern for matching files in ``input_dir_background``", + ).tag(config=True) + + dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( + config=True + ) + + stack_telescope_images = Bool( + default_value=False, + allow_none=False, + help=( + "Set whether to stack the telescope images in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + sort_by_intensity = Bool( + default_value=False, + allow_none=True, + help=( + "Set whether to sort the telescope images by intensity in the data loader. " + "Requires DLDataReader mode to be ``stereo``." + ), + ).tag(config=True) + + output_dir = Path( + exits=False, + default_value=None, + allow_none=False, + directory_ok=True, + file_ok=False, + help="Output directory for the trained reconstructor.", + ).tag(config=True) + + reco_tasks = List( + trait=CaselessStrEnum(["type", "energy", "cameradirection", "skydirection"]), + allow_none=False, + default_value=None, + help=( + "List of reconstruction tasks to perform. " + "'type': classification of the primary particle type " + "'energy': regression of the primary particle energy " + "'cameradirection': regression of the primary particle arrival direction in camera coordinates " + "'skydirection': regression of the primary particle arrival direction in sky coordinates" + ), + ).tag(config=True) + + n_epochs = Int( + default_value=10, + allow_none=False, + help="Number of epochs to train the neural network.", + ).tag(config=True) + + batch_size = Int( + default_value=64, + allow_none=False, + help="Size of the batch to train the neural network.", + ).tag(config=True) + + validation_split = Float( + default_value=0.1, + help="Fraction of the data to use for validation", + min=0.01, + max=0.99, + ).tag(config=True) + + optimizer = Dict( + default_value={ + "name": "Adam", + "base_learning_rate": 0.0001, + "adam_epsilon": 1.0e-8, + }, + help=( + "Optimizer to use for training. " + "E.g. {'name': 'Adam', 'base_learning_rate': 0.0001, 'adam_epsilon': 1.0e-8}. " + ), + ).tag(config=True) + + random_seed = Int( + default_value=0, + help=( + "Random seed for shuffling the data " + "before the training/validation split " + "and after the end of an epoch." + ), + ).tag(config=True) + + save_onnx = Bool( + default_value=False, + allow_none=False, + help="Set whether to save model in an ONNX file.", + ).tag(config=True) + + overwrite = Bool(help="Overwrite output dir if it exists").tag(config=True) + + # Unified Hardware Architecture & Execution Strategy + device = CaselessStrEnum( + ["cuda", "cpu", "mps"], + default_value="cuda", + help="Device to use for training: cuda, cpu, or mps", + ).tag(config=True) + + precision_type = CaselessStrEnum( + ["64-true", "32-true", "16-true", "16-mixed", "bf16-mixed", "bf16-true"], + default_value="32-true", + help="Precision for type classification model", + ).tag(config=True) + + precision_energy = CaselessStrEnum( + ["64-true", "32-true", "16-true", "16-mixed", "bf16-mixed", "bf16-true"], + default_value="32-true", + help="Precision for energy regression model", + ).tag(config=True) + + precision_direction = CaselessStrEnum( + ["64-true", "32-true", "16-true", "16-mixed", "bf16-mixed", "bf16-true"], + default_value="32-true", + help="Precision for direction regression model", + ).tag(config=True) + + devices = List( + trait=Int(), + default_value=[0], + help="List of GPU device IDs to use for training", + ).tag(config=True) + + strategy = Unicode( + default_value="auto", + help="Multi-GPU strategy (e.g. auto, ddp, deepspeed_stage_2, fsdp)", + ).tag(config=True) + + # Optimizer & Learning Rate Hyperparameters + optimizer_momentum = Float( + default_value=0.9, + help="Momentum for optimizer (e.g. SGD, RMSProp, Adamw)", + ).tag(config=True) + + optimizer_weight_decay = Float( + default_value=0.0005, + help="Weight decay for optimizer", + ).tag(config=True) + + lrf = Float( + default_value=0.1, + help="Learning rate factor for learning rate schedulers", + ).tag(config=True) + + l2_lambda = Float( + default_value=1e-7, + help="L2 regularization lambda factor", + ).tag(config=True) + + gradient_clip_val = Float( + default_value=3.0, + allow_none=True, + help="Gradient clipping value to avoid gradient explosion", + ).tag(config=True) + + save_k_checkpoints = Int( + default_value=3, + help="Number of top checkpoints to save", + ).tag(config=True) + + experiment_number = Int( + default_value=1, + help="Experiment run number", + ).tag(config=True) + + # Event Cut-offs + leakage_intensity_cutoff = Float( + default_value=0.2, + help="Events with leakage intensity greater than this value are removed", + ).tag(config=True) + + intensity_cutoff = Float( + default_value=50.0, + help="Events with Hillas intensity below this value are removed", + ).tag(config=True) + + # Normalizations + apply_log_scaling = List( + trait=Bool(), + default_value=[True, True], + help="List specifying whether to apply log10(X+1.0) scaling for [charge, peak_time]", + ).tag(config=True) + + use_clean = Bool( + default_value=True, + help="Use the image with the applied mask (True), otherwise the mask is not applied", + ).tag(config=True) + + use_clean_dvr = Bool( + default_value=False, + help="Use clean DVR mask", + ).tag(config=True) + + type_mu = Float( + default_value=0.0, + help="Mean for type channel normalization", + ).tag(config=True) + + type_sigma = Float( + default_value=1.0, + help="Std dev for type channel normalization", + ).tag(config=True) + + dir_mu = Float( + default_value=0.0, + help="Mean for direction channel normalization", + ).tag(config=True) + + dir_sigma = Float( + default_value=1.0, + help="Std dev for direction channel normalization", + ).tag(config=True) + + energy_mu = Float( + default_value=0.0, + help="Mean for energy channel normalization", + ).tag(config=True) + + energy_sigma = Float( + default_value=1.0, + help="Std dev for energy channel normalization", + ).tag(config=True) + + # Augmentations + use_augmentation = Bool( + default_value=False, + help="Apply data augmentation during training", + ).tag(config=True) + + aug_prob = Float( + default_value=0.5, + help="Probability of applying data augmentation", + ).tag(config=True) + + rot_prob = Float( + default_value=0.5, + help="Probability of applying rotation", + ).tag(config=True) + + trans_prob = Float( + default_value=0.5, + help="Probability of applying translation", + ).tag(config=True) + + flip_hor_prob = Float( + default_value=0.5, + help="Probability of applying horizontal flip", + ).tag(config=True) + + flip_ver_prob = Float( + default_value=0.5, + help="Probability of applying vertical flip", + ).tag(config=True) + + mask_prob = Float( + default_value=0.5, + help="Probability of applying masking", + ).tag(config=True) + + mask_dvr_prob = Float( + default_value=0.5, + help="Probability of applying DVR masking", + ).tag(config=True) + + noise_prob = Float( + default_value=0.5, + help="Probability of applying noise", + ).tag(config=True) + + max_rot = Float( + default_value=5.0, + help="Maximum rotation in augmentation (degrees)", + ).tag(config=True) + + max_trans = Float( + default_value=10.0, + help="Maximum translation in augmentation (pixels)", + ).tag(config=True) + + # Dataset & Loading Settings + num_workers = Int( + default_value=1, + help="Number of workers for data loading", + ).tag(config=True) + + pin_memory = Bool( + default_value=True, + help="Pin memory in PyTorch dataloader for faster GPU transfer", + ).tag(config=True) + + persistent_workers = Bool( + default_value=True, + help="Keep workers alive between epochs", + ).tag(config=True) + + # Checkpoint paths + type_checkpoint = Path( + default_value=None, + allow_none=True, + help="Path to checkpoint for type classification", + ).tag(config=True) + + energy_checkpoint = Path( + default_value=None, + allow_none=True, + help="Path to checkpoint for energy regression", + ).tag(config=True) + + direction_checkpoint = Path( + default_value=None, + allow_none=True, + help="Path to checkpoint for direction regression", + ).tag(config=True) + + load_onnx_model = Path( + default_value=None, + allow_none=True, + exists=True, + help="Path to an ONNX model to load and train/fine-tune.", + ).tag(config=True) + + aliases = { + "framework": "TrainCTLearnModel.framework_type", + "n_epochs": "TrainCTLearnModel.n_epochs", + "signal": "TrainCTLearnModel.input_dir_signal", + "background": "TrainCTLearnModel.input_dir_background", + "pattern-signal": "TrainCTLearnModel.file_pattern_signal", + "pattern-background": "TrainCTLearnModel.file_pattern_background", + "reco": "TrainCTLearnModel.reco_tasks", + "save_onnx": "TrainCTLearnModel.save_onnx", + "load_onnx_model": "TrainCTLearnModel.load_onnx_model", + "experiment_number": "TrainCTLearnModel.experiment_number", + "apply_log_scaling": "TrainCTLearnModel.apply_log_scaling", + "use_clean": "TrainCTLearnModel.use_clean", + "random_seed": "TrainCTLearnModel.random_seed", + "optimizer": "TrainCTLearnModel.optimizer", + "overwrite": "TrainCTLearnModel.overwrite", + "validation_split": "TrainCTLearnModel.validation_split", + "batch_size": "TrainCTLearnModel.batch_size", + "sort_by_intensity": "TrainCTLearnModel.sort_by_intensity", + "stack_telescope_images": "TrainCTLearnModel.stack_telescope_images", + "dl1dh_reader_type": "TrainCTLearnModel.dl1dh_reader_type", + ("o", "output"): "TrainCTLearnModel.output_dir", + "device": "TrainCTLearnModel.device", + "devices": "TrainCTLearnModel.devices", + "strategy": "TrainCTLearnModel.strategy", + "use_augmentation": "TrainCTLearnModel.use_augmentation", + "precision_type": "TrainCTLearnModel.precision_type", + "precision_energy": "TrainCTLearnModel.precision_energy", + "precision_direction": "TrainCTLearnModel.precision_direction", + "type_checkpoint": "TrainCTLearnModel.type_checkpoint", + "energy_checkpoint": "TrainCTLearnModel.energy_checkpoint", + "direction_checkpoint": "TrainCTLearnModel.direction_checkpoint", + "num_workers": "TrainCTLearnModel.num_workers", + "pin_memory": "TrainCTLearnModel.pin_memory", + "persistent_workers": "TrainCTLearnModel.persistent_workers", + } + + flags = { + "overwrite": ( + {"TrainCTLearnModel": {"overwrite": True}}, + "Overwrite existing files", + ), + } + + classes = ( + [ + DLDataReader, + ] + + classes_with_traits(DLDataReader) + ) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def setup(self): + + + # Check if the output directory exists and if it should be overwritten + if self.output_dir.exists(): + if not self.overwrite: + raise ToolConfigurationError( + f"Output directory {self.output_dir} already exists. Use --overwrite to overwrite." + ) + else: + # Remove the output directory if it exists + self.log.info("Removing existing output directory %s", self.output_dir) + shutil.rmtree(self.output_dir) + + # Get signal input files + self.input_url_signal = [] + for signal_pattern in self.file_pattern_signal: + self.input_url_signal.extend(self.input_dir_signal.glob(signal_pattern)) + + # Get bkg input files + self.input_url_background = [] + if self.input_dir_background is not None: + for background_pattern in self.file_pattern_background: + self.input_url_background.extend( + self.input_dir_background.glob(background_pattern) + ) + + # Set up the data reader + self.log.info("Loading data:") + self.log.info("For a large dataset, this may take a while...") + if self.dl1dh_reader_type == "DLFeatureVectorReader": + raise NotImplementedError( + "'DLFeatureVectorReader' is not supported in CTLearn yet. " + "Missing stereo CTLearnModel implementation." + ) + if "quality_criteria" not in self.config.get("TableQualityQuery", {}): + self.quality_query = TableQualityQuery(quality_criteria=[('> 50 phe', 'hillas_intensity > 50')]) + self.config.TableQualityQuery.quality_criteria = self.quality_query.quality_criteria + + if self.dl1dh_reader_type == "DLImageReader": + self.channels = ["cleaned_image", "cleaned_peak_time"] + if self.apply_log_scaling[0]: + self.channels[0] = "log_" + self.channels[0] + if "channels" not in self.config.get("DLImageReader", {}): + self.config.DLImageReader.channels = self.channels + + print(f"self.dl1dh_reader_type: {self.dl1dh_reader_type}") + self.dl1dh_reader = DLDataReader.from_name( + self.dl1dh_reader_type, + input_url_signal=sorted(self.input_url_signal), + input_url_background=sorted(self.input_url_background), + parent=self + ) + + # self.dl1dh_reader.channels = ["cleaned_image", "cleaned_peak_time"] + # self.dl1dh_reader.quality_query.quality_criteria = [('> 50 phe', 'hillas_intensity > 1000')] + + self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) + if "type" in self.reco_tasks: + self.log.info( + "Number of signal events: %d", self.dl1dh_reader.n_signal_events + ) + self.log.info( + "Number of background events: %d", self.dl1dh_reader.n_bkg_events + ) + # Check if the number of events is enough to form a batch + if self.dl1dh_reader._get_n_events() < self.batch_size: + raise ValueError( + f"{self.dl1dh_reader._get_n_events()} events are not enough " + f"to form a batch of size {self.batch_size}. Reduce the batch size." + ) + + self.log.info(f'Class weight{self.dl1dh_reader.class_weight}') + + # Check if there are at least two classes in the reader for the particle classification + if "type" in self.reco_tasks: + if self.dl1dh_reader.class_weight is None and "type" in self.reco_tasks: + raise ValueError( + "Classification task selected but less than two classes are present in the data." + ) + if self.dl1dh_reader.class_weight[0]/self.dl1dh_reader.class_weight[1] > 1.1: + self.log.warning( + "The dataset seems to be imbalanced. " + f"Consider using more background events using class_weight in the optimizer.: {self.dl1dh_reader.class_weight}" + ) + if self.dl1dh_reader.class_weight[0]/self.dl1dh_reader.class_weight[1] < 0.9: + self.log.warning( + "The dataset seems to be imbalanced. " + f"Consider using more signal events using class_weight in the optimizer.: {self.dl1dh_reader.class_weight}" + ) + # Check if stereo mode is selected for stacking telescope images + if self.stack_telescope_images and self.dl1dh_reader.mode == "mono": + raise ToolConfigurationError( + f"Cannot stack telescope images in mono mode. Use stereo mode for stacking." + ) + # Ckeck if only one telescope type is selected for stacking telescope images + if ( + self.stack_telescope_images + and len(list(self.dl1dh_reader.selected_telescopes)) > 1 + ): + raise ToolConfigurationError( + f"Cannot stack telescope images from multiple telescope types. Use only one telescope type." + ) + # Check if sorting by intensity is disabled for stacking telescope images + if self.stack_telescope_images and self.sort_by_intensity: + raise ToolConfigurationError( + f"Cannot stack telescope images when sorting by intensity. Disable sorting by intensity." + ) + + def start(self): + pass + + def finish(self): + pass diff --git a/ctlearn/tools/train/conftest.py b/ctlearn/tools/train/conftest.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/tools/train/keras/__init__.py b/ctlearn/tools/train/keras/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/tools/train/keras/train_keras_model.py b/ctlearn/tools/train/keras/train_keras_model.py new file mode 100644 index 00000000..af09b2fe --- /dev/null +++ b/ctlearn/tools/train/keras/train_keras_model.py @@ -0,0 +1,362 @@ +import atexit +import pandas as pd +import numpy as np +import shutil +import tensorflow as tf +from tqdm import tqdm +from time import time +from keras.callbacks import Callback + +from ctapipe.core.traits import ( + Bool, + CaselessStrEnum, + Path, + Float, + Int, + List, + Dict, + classes_with_traits, + ComponentName, + Unicode, +) +from ctlearn.core.data_loader.loader import DLDataLoader +from ctlearn.tools.train.base_train_model import TrainCTLearnModel +from ctlearn.core.model import CTLearnModel +from ctlearn.utils import validate_trait_dict + +try: + import keras +except ImportError: + raise ImportError("keras is not installed in your environment!") + +class TqdmProgressBar(Callback): + def on_epoch_begin(self, epoch, logs=None): + self.epoch_start_time = time() + self.progress_bar = tqdm(total=self.params['steps'], desc=f'Epoc {epoch + 1}/{self.params["epochs"]}', unit='batch') + + def on_epoch_end(self, epoch, logs=None): + epoch_duration = time() - self.epoch_start_time + print(f'\nEpoch Time {epoch + 1}: {epoch_duration:.2f} seconds') + self.progress_bar.close() + + def on_batch_begin(self, batch, logs=None): + self.batch_start_time = time() + + def on_batch_end(self, batch, logs=None): + # batch_duration = time() - self.batch_start_time + self.progress_bar.set_postfix(loss=logs.get('loss'), val_loss=logs.get('val_loss')) + self.progress_bar.update(1) + +class TrainKerasModel(TrainCTLearnModel): + """ + Tool to train a ``~ctlearn.core.model.CTLearnModel`` on R1/DL1a data using keras. + + The tool sets up the keras model using the specified optimizer and callbacks. The keras model is trained + on the input data (R1 calibrated waveforms or DL1a images) and saved in the output directory. + """ + + name = "ctlearn-train-keras-model" + description = __doc__ + + examples = """ + To train a 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 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 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 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 \\ + """ + + model_type = ComponentName( + CTLearnModel, default_value="KerasResNet" + ).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) + + 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) + + early_stopping = Dict( + default_value=None, + allow_none=True, + help=( + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " + ) + ).tag(config=True) + + aliases = { + **TrainCTLearnModel.aliases, + } + + def setup(self): + + print(tf.config.list_physical_devices('GPU')) + # Create a MirroredStrategy. + self._keras_strategy = tf.distribute.MirroredStrategy() + atexit.register(self._keras_strategy._extended._collective_ops._lock.locked) # type: ignore + self.log.info("Number of devices: %s", self._keras_strategy.num_replicas_in_sync) + # print(self.framework_type) + super().setup() + + # Set up the data loaders for training and validation + indices = list(range(self.dl1dh_reader._get_n_events())) + # Shuffle the indices before the training/validation split + np.random.seed(self.random_seed) + np.random.shuffle(indices) + n_validation_examples = int( + self.validation_split * self.dl1dh_reader._get_n_events() + ) + training_indices = indices[n_validation_examples:] + validation_indices = indices[:n_validation_examples] + + # Set self.strategy.num_replicas_in_sync to 1 in case that does not exist (Pytorch) + if not hasattr(self, "_keras_strategy"): + self._keras_strategy = type("FakeStrategy", (), {"num_replicas_in_sync": 1})() + print("num_replicas_in_sync:", self._keras_strategy.num_replicas_in_sync) + + print("BASE TRAIN FRAMEWORK", self.framework_type) + + self.training_loader = DLDataLoader.create( + framework=self.framework_type, + DLDataReader=self.dl1dh_reader, + indices=training_indices, + tasks=self.reco_tasks, + batch_size=self.batch_size * self._keras_strategy.num_replicas_in_sync, + random_seed=self.random_seed, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + + self.validation_loader = DLDataLoader.create( + framework=self.framework_type, + DLDataReader=self.dl1dh_reader, + indices=validation_indices, + tasks=self.reco_tasks, + batch_size=self.batch_size * self._keras_strategy.num_replicas_in_sync, + random_seed=self.random_seed, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + ) + + def start(self): + print("Start KERAS") + print("EPOCHS:",self.n_epochs) + print("save_onnx:",self.save_onnx) + # Set up the keras callbacks + monitor = "val_loss" + monitor_mode = "min" + # Model checkpoint callback + # Temp fix for supporting keras2 & keras3 + if int(keras.__version__.split(".")[0]) >= 3: + model_path = f"{self.output_dir}/ctlearn_model.keras" + else: + model_path = f"{self.output_dir}/ctlearn_model.cpk" + model_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 + 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) + # Open a strategy scope. + with self._keras_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) + + tqdm_callback = TqdmProgressBar() + self.callbacks.append(tqdm_callback) + + # 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 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 \ No newline at end of file diff --git a/ctlearn/tools/train/pytorch/CTLearnPL.py b/ctlearn/tools/train/pytorch/CTLearnPL.py new file mode 100644 index 00000000..cbb19519 --- /dev/null +++ b/ctlearn/tools/train/pytorch/CTLearnPL.py @@ -0,0 +1,1770 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +import numpy as np +from ctlearn.core.pytorch.nets.loss_functions.loss_functions import ( + FocalLoss, + VectorLoss, + AngularDistance, + cosine_direction_loss, +) +from ctlearn.core.pytorch.nets.optimizer.optimizer import one_cycle +from tqdm import tqdm +import os +import math +import torch.optim.lr_scheduler as lr_scheduler +from sklearn.metrics import confusion_matrix +from torchmetrics.classification import ConfusionMatrix, MulticlassPrecision, MulticlassF1Score + +from io import BytesIO +import torch.nn.functional as F +from ctlearn.core.pytorch.utils import utils +from ctlearn.core.pytorch.visualization.vis_utils import ( + plot_confusion_matrix, + plot_energy_resolution_error, + plot_direction_resolution_error, +) + +# import utils_torch +import pytorch_lightning as pl +from torchmetrics import Accuracy +import matplotlib.pyplot as plt +from astropy.coordinates import SkyCoord +from astropy.time import Time +import astropy.units as u +from ctlearn.core.ctlearn_enum import Task, Mode +import gc +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from ctlearn.core.pytorch.nets.loss_functions.loss_functions import evidential_regression_loss +import inspect + +import matplotlib +matplotlib.use('Agg') +torch.cuda.empty_cache() +os.environ["CUDA_LAUNCH_BLOCKING"] = "1" + +class CTLearnTrainer(pl.Trainer): + def __init__(self,**kwargs): + super(CTLearnTrainer, self).__init__(**kwargs) + self.multi_gpu=False + if self.world_size > 1: + self.multi_gpu=True + + def predict( + self, + model, + dataloaders=None, + return_predictions=False, + h5_file_name="./results.r1.dl2.h5", + task: Task = None, + mode: Mode = None, + **kwargs, + ): + + # ---------------------------------------------------------------- + # Prediction + # ---------------------------------------------------------------- + self.class_predictions = [] + self.energy_predictions = [] + self.direction_predictions = [] + self.event_id_list = [] + self.obs_id_list = [] + self.labels_energy_list = [] + self.labels_direction_list = [] + self.labels_true_alt_az_list = [] + self.hillas_list = [] + self.pointing_dir = {} + self.h5_file_name = h5_file_name + self.task = task + self.mode = mode + # super(CTLearnTrainer, self).__init__(kwargs) + print("Using CustomTrainer's predict method.") + + if dataloaders is None: + raise ValueError("A dataloader must be provided for prediction.") + + model.eval() + with torch.no_grad(): + # Call your custom logic here + results = model.generate_results( + input_data_loader=dataloaders, + h5_file_name=None, + task=task, + mode=mode, + ) + + if return_predictions: + return results + + print("Prediction complete.") + + def get_log_dir(self) -> str: + return self.logger.log_dir + +class CTLearnPL(pl.LightningModule): + + # def setup(self,stage): + + # # self.dummy = None + + # self.dummy = self.occupy_free_gpu_memory(self.device) + + def occupy_free_gpu_memory(self,device): + # Get free and total GPU memory using CUDA APIs + free_mem, total_mem = torch.cuda.mem_get_info(device) + + # Leave a margin (e.g., 100 MB) + margin = 20000 * 1024 * 1024 # 1000 MB + mem_to_allocate = max(0, free_mem - margin) + + if mem_to_allocate > 0: + try: + # Each float32 element takes 4 bytes + num_elements = mem_to_allocate // 4 + dummy_tensor = torch.zeros(num_elements, dtype=torch.float32, device=device) + print(f"Occupied {mem_to_allocate/1024**2:.2f} MB of GPU memory on {device}") + return dummy_tensor + except RuntimeError as e: + print(f"Could not allocate memory: {e}") + return None + else: + print("Not enough free GPU memory to occupy.") + return None + + def __init__( + self, + model, + save_folder, + task: Task, + mode: Mode, + parameters, + train_loader=None, + val_loader=None, + test_val_loader=None, + train_dataset=None, + val_dataset=None, + num_inputs=1, + k=3, + + ): + super(CTLearnPL, self).__init__() + + self.task = task + self.mode = mode + + self.save_folder = save_folder + self.model = model + + # torch.autograd.set_detect_anomaly(True) + + self._device = torch.device(parameters["arch"]["device"]) + self.device_type = parameters["arch"]["device"] + + self.model.to(self.device) + + + + # # Ejecutar la función + # self.dummy = None + + # self.occupy_free_gpu_memory(self.device) + self.k = k # Number of top results to save + # Get the number of inputs of the net + sig = inspect.signature(model.forward) + num_inputs = len(sig.parameters) + # Detect if model is diffusive + if hasattr(self.model,"T"): + self.is_difussion=True + else: + self.is_difussion=False + print("Diffusion: ",self.is_difussion) + + self.num_inputs = num_inputs + self.train_loader = train_loader + self.val_loader = val_loader + self.train_dataset = train_dataset + self.val_dataset = val_dataset + self.test_val_loader = test_val_loader + self.class_names = ["gamma", "proton"] + self.val_proton_gammanes = [] + self.val_gamma_gammanes = [] + # Hyperparameters + self.set_hyperparameters(parameters) + + self.optimizer = None + self.scheduler = None + # self.class_weights = torch.tensor([1.0, 1.3], dtype=torch.float32).to(self.device).contiguous() + + # Loss Function + if self.task == Task.type: + weights = ( + torch.tensor([parameters['class_weight'][0],parameters['class_weight'][1]], dtype=torch.float32).to(self.device).contiguous() + ) # [1.0, 1.3] + self.register_buffer('class_weights', weights) + self.criterion_class = nn.CrossEntropyLoss( + weight=self.class_weights, reduction="mean" + ) + + + self.criterion_energy_value = torch.nn.L1Loss(reduction="sum") + self.criterion_direction = torch.nn.SmoothL1Loss() # nn.MSELoss() + self.criterion_magnitud = torch.nn.L1Loss(reduction="mean") + + self.criterion_direction_none = torch.nn.SmoothL1Loss(reduction="none") # nn.MSELoss() + self.criterion_magnitud_none = torch.nn.L1Loss(reduction="none") + self.criterion_alt_az_l1_none = torch.nn.L1Loss(reduction="none") + self.criterion_energy_value_none = torch.nn.L1Loss(reduction="none") + + self.criterion_direction = torch.nn.L1Loss(reduction="mean") # nn.MSELoss() + self.criterion_vector = VectorLoss(alpha=0.1, reduction="mean") + self.criterion_alt_az_l1 = torch.nn.L1Loss(reduction="mean") + + self.criterion_alt_az = evidential_regression_loss(lamb=0.01, reduction="mean") + + + self.best_loss = float("inf") + self.best_accuracy = 0 + # Best Metrics Tracking + self.best_losses = [(float("inf"), None)] * k + self.best_accuracies = [(0, None)] * k + self.best_validation_accuracy = 0 + + self.correct_classification = 0 + + self.f1_score_val = MulticlassF1Score( + num_classes=2, + dist_sync_on_step=True, + ) + + self.f1_score_train = MulticlassF1Score( + num_classes=2, + dist_sync_on_step=True, + ) + + self.precision_val = MulticlassPrecision(num_classes=2,dist_sync_on_step=True,) + self.precision_train = MulticlassPrecision(num_classes=2,dist_sync_on_step=True,) + + self.class_train_accuracy = Accuracy( + task="multiclass", + num_classes=2, + dist_sync_on_step=True # GPUs Sync + ) + + self.class_val_accuracy = Accuracy( + task="multiclass", + num_classes=2, + dist_sync_on_step=True # GPUs Sync + ) + + self.confusion_matrix = ConfusionMatrix(num_classes=2, task="multiclass",dist_sync_on_step=True) + + self.loss_train_sum = 0.0 + self.num_train_batches = 0 + + # ---------------------------------------------------------------- + # List used to plot on tensorboard + # ---------------------------------------------------------------- + + # ---------------------------------------------------------------- + # Validation + # ---------------------------------------------------------------- + self.loss_val_sum = 0.0 + self.num_val_batches = 0 + + + self.val_angular_diff_list = [] + self.val_energy_diff_list = [] + + self.val_energy_pred_list = [] + + self.val_alt_pred_list = [] + self.val_az_pred_list = [] + self.val_alt_label_list = [] + self.val_az_label_list = [] + + self.loss_val_distance = 0.0 + self.loss_val_dx_dy = 0.0 + self.loss_val_angular_error = 0.0 + + self.val_energy_label_list = [] + self.val_hillas_intensity_list = [] + + # ---------------------------------------------------------------- + # Training + # ---------------------------------------------------------------- + self.loss_train_distance = 0.0 + self.loss_train_dx_dy = 0.0 + self.loss_train_angular_error = 0.0 + + self.alt_off_list = [] + self.az_off_list = [] + + # ---------------------------------------------------------------- + # Predictions + # ---------------------------------------------------------------- + + self.predictions = [] + + self.class_predictions = [] + self.energy_predictions = [] + self.direction_predictions = [] + self.event_id_list = [] + self.obs_id_list = [] + self.labels_energy_list = [] + self.labels_direction_list = [] + self.labels_true_alt_az_list = [] + self.hillas_list = [] + self.pointing = [] + # ---------------------------------------------------------------------------------------------------------- + def train_dataloader(self): + return self.train_loader + # ---------------------------------------------------------------------------------------------------------- + def set_hyperparameters(self, parameters): + # Hyperparameters + self.learning_rate = float(parameters["hyp"]["learning_rate"]) + self.adam_epsilon = float(parameters["hyp"]["adam_epsilon"]) + self.momentum = float(parameters["hyp"]["momentum"]) + self.weight_decay = float(parameters["hyp"]["weight_decay"]) + self.num_epochs = int(parameters["hyp"]["epochs"]) + + self.start_epoch = parameters["hyp"]["start_epoch"] + self.steps_epoch = parameters["hyp"]["steps_epoch"] + self.lrf = float(parameters["hyp"]["lrf"]) + self.optimizer_type = str(parameters["hyp"]["optimizer"]).lower() + self.l2_lambda = float(parameters["hyp"]["l2_lambda"]) + # ---------------------------------------------------------------------------------------------------------- + def forward(self, x): + return self.model(x) + # ---------------------------------------------------------------------------------------------------------- + def save_checkpoint(self, save_folder, metric_value, filename_prefix, is_loss=True): + """Save checkpoint and manage top k checkpoints for loss or accuracy.""" + filename = os.path.join( + save_folder, f"{filename_prefix}_{metric_value:.16f}.pth" + ) + torch.save( + { + "model_state_dict": self.model.state_dict(), + "optimizer_state_dict": self.optimizer.state_dict(), + "metric_value": metric_value, + }, + filename, + ) + + # Determine the list to update + current_list = self.best_losses if is_loss else self.best_accuracies + current_list.append((metric_value, filename)) + current_list.sort(reverse=not is_loss) + if len(current_list) > self.k: + removed_metric, removed_file = current_list.pop() + if removed_file and os.path.exists(removed_file): + # TODO: Add a Lock to avoid problems when using multiple GPUS + try: + os.remove(removed_file) # Remove the worst performing file + except FileNotFoundError: + pass # File doesn't exist, continue execution + + + # ---------------------------------------------------------------------------------------------------------- + def compute_type_loss( + self, classification_pred, labels_class, test_val=False, training=False + ): + target = labels_class.to(torch.int64) + + # Cálculo de la loss con F.cross_entropy + loss_class = F.cross_entropy(classification_pred, target, weight=self.class_weights.to(classification_pred.device), reduction='mean') + + # Calculate accuracy + predicted = torch.softmax(classification_pred, dim=1) + predicted = predicted.argmax(dim=1) + + + accuracy = 0 + precision = 0 + loss = loss_class + + if training: + + self.class_train_accuracy.update(predicted, labels_class) + accuracy = self.class_train_accuracy.compute().item() + self.f1_score_train.update(predicted, labels_class) + self.precision_train.update(predicted, labels_class) + precision = self.precision_train.compute().item() + else: + + self.class_val_accuracy.update(predicted, labels_class) + accuracy = self.class_val_accuracy.compute().item() + self.confusion_matrix.update(predicted, labels_class) + self.f1_score_val.update(predicted, labels_class) + self.precision_val.update(predicted, labels_class) + precision = self.precision_val.compute().item() + + return loss, accuracy, predicted, precision + # ---------------------------------------------------------------------------------------------------------- + def compute_camera_direction_loss(self, direction_pred, labels_direction, training=False): + + labels_dx_dy = labels_direction[:, 0:2] + label_distance = labels_direction[:, 2] + + + if isinstance(direction_pred, tuple): + direction_pred = list(direction_pred) + + pred_dx_dy = direction_pred[0][:,0:2].unsqueeze(-1) + pred_distance = direction_pred[0][:,2].unsqueeze(-1) + else: + pred_dx_dy = direction_pred[:,0:2].unsqueeze(-1) + pred_distance = direction_pred[:,2].unsqueeze(-1) + + + loss_angular_diff = cosine_direction_loss(pred_dx_dy[:,0],pred_dx_dy[:,1], labels_dx_dy[:, 0],labels_dx_dy[:, 1]) + + + _, angular_diff = AngularDistance( + pred_dx_dy[:,0], + labels_dx_dy[:, 0], + pred_dx_dy[:,1], + labels_dx_dy[:, 1], + reduction="None", + ) + + vector_cam_distance = torch.sqrt(pred_dx_dy[:,0]**2 + pred_dx_dy[:,1]**2) + loss_dx_dy = self.criterion_alt_az_l1(pred_dx_dy, labels_dx_dy) + loss_distance = self.criterion_magnitud(label_distance, pred_distance) + loss_distance_dx_dy = self.criterion_magnitud(label_distance, vector_cam_distance) + + loss = loss_dx_dy + loss_distance + loss_distance_dx_dy + loss_angular_diff + return loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_energy_loss( + self, energy_pred, labels_energy, test_val=False, training=False + ): + + loss_energy = self.criterion_energy_value(energy_pred, labels_energy) + loss = loss_energy + #----------------------------------------- + # k=0.6 + # e_thrs= -0.3 + # loss_energy = self.criterion_energy_value_none(energy_pred, labels_energy) + # energy_tev = torch.pow(10,labels_energy) + # energy_weight=k*torch.exp(1+(np.log10(1/k)*(energy_tev-e_thrs))) + # weight_loss = energy_weight * (loss_energy) + + # loss = weight_loss.sum() + #----------------------------------------- + + if training == False: + energy_pred = pow(10, energy_pred) + labels_energy = pow(10, labels_energy) + energy_diff = torch.abs(energy_pred - labels_energy) + energy_diff = energy_diff.float().cpu().detach().numpy() + else: + energy_diff = None + + return loss, energy_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_energy_loss_diffusion(self, x_1, y,training=False, x_2 = None + # self, energy_pred, labels_energy, test_val=False, training=False + ): + + loss = 0 + # y = y.squeeze(-1) + y_embed = self.model.target_embedder(y) + k=0.6 + e_thrs= -0.3 + for t in range(self.model.T): + # Add noise to target (landmarks) + alpha_bar_t = self.model.alpha_bar[t] + noise = torch.randn_like(y_embed) + z_t = torch.sqrt(alpha_bar_t) * y_embed + torch.sqrt(1 - alpha_bar_t) * noise + + if x_2 is None: + # Denoise step + z, _ = self.model.blocks[t](x_1, z_t, None) # W_embed not needed + else: + z, _ = self.model.blocks[t](x_1 ,x_2, z_t, None) # W_embed not needed + + + # Loss to clean target + loss_l2 = F.mse_loss(z, y_embed) + + labels_energy = y + + # Weighted by SNR difference + step_loss = 2.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + # Final step: use regression head + if t == self.model.T - 1: + + energy_pred = self.model.regress(z) + + # labels_energy_tev = torch.pow(10,labels_energy) + # energy_pred_tev = torch.pow(10,energy_pred) + + # loss_energy = self.criterion_energy_value(energy_pred_tev, labels_energy_tev) + + loss_energy = self.criterion_energy_value(energy_pred, labels_energy) + # loss_energy = F.mse_loss(energy_pred, labels_energy,reduction="sum") + # loss_energy = self.criterion_energy_value_none(energy_pred, labels_energy) + + # energy_weight=k*torch.exp(1+(np.log10(1/k)*(energy_tev-e_thrs))) + # weight_loss = energy_weight * (loss_energy) + + # step_loss = step_loss + weight_loss.mean() + + if training == False: + energy_pred = pow(10, energy_pred) + labels_energy = pow(10, labels_energy) + energy_diff = torch.abs(energy_pred - labels_energy) + energy_diff = energy_diff.float().cpu().detach().numpy() + else: + energy_diff = None + + step_loss = step_loss + loss_energy + + loss = loss + step_loss + + return loss, energy_diff, energy_pred + # ---------------------------------------------------------------------------------------------------------- + def compute_camera_direction_loss_diffusion(self, x_1, y, labels_energy_value, x_2 = None): + + loss = 0 + y = y.squeeze(-1) + y_embed = self.model.target_embedder(y) + k=0.6 + e_thrs= -0.3 + + for t in range(self.model.T): + # Add noise to target (landmarks) + alpha_bar_t = self.model.alpha_bar[t] + noise = torch.randn_like(y_embed) + z_t = torch.sqrt(alpha_bar_t) * y_embed + torch.sqrt(1 - alpha_bar_t) * noise + + if x_2 is None: + # Denoise step + z, _ = self.model.blocks[t](x_1, z_t, None) # W_embed not needed + else: + z, _ = self.model.blocks[t](x_1,x_2, z_t, None) # W_embed not needed + + preds = self.model.regress(z) + # Loss to clean target + loss_l2 = F.mse_loss(preds, y) + + # Weighted by SNR difference + step_loss = 2.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + # Final step: use regression head + if t == self.model.T - 1: + labels_dx_dy = y[:, 0:2] + label_distance = y[:, 2] + direction_pred = self.model.regress(z) + if isinstance(direction_pred, tuple): + # Not Tested + direction_pred = list(direction_pred) + + pred_dx_dy = direction_pred[0][:,0:2].unsqueeze(-1) + pred_distance = direction_pred[0][:,2].unsqueeze(-1) + else: + pred_dx_dy = direction_pred[:,0:2] + pred_distance = direction_pred[:,2] + + # loss_angular_diff = cosine_direction_loss(pred_dx_dy[:,0],pred_dx_dy[:,1], labels_dx_dy[:, 0],labels_dx_dy[:, 1]) + loss_angular_diff = cosine_direction_loss(pred_dx_dy[:,0],pred_dx_dy[:,1], labels_dx_dy[:, 0],labels_dx_dy[:, 1],reduction="none") + _, angular_diff = AngularDistance( + pred_dx_dy[:,0], + labels_dx_dy[:, 0], + pred_dx_dy[:,1], + labels_dx_dy[:, 1], + reduction="None", + + ) + + vector_cam_distance = torch.sqrt(pred_dx_dy[:,0]**2 + pred_dx_dy[:,1]**2) + # loss_dx_dy = self.criterion_alt_az_l1(pred_dx_dy, labels_dx_dy) + # loss_distance = self.criterion_magnitud(label_distance, pred_distance) + # loss_distance_dx_dy = self.criterion_magnitud(label_distance, vector_cam_distance) + loss_dx_dy = self.criterion_alt_az_l1_none(pred_dx_dy, labels_dx_dy).mean(dim=1) + loss_distance = self.criterion_magnitud_none(label_distance, pred_distance) + loss_distance_dx_dy = self.criterion_magnitud_none(label_distance, vector_cam_distance) + + energy = torch.pow(10,labels_energy_value) + # k=4.3 + # e_thrs= 4 + # energy_weight = k*(1/(1+torch.exp(-(1/k)*(energy-e_thrs)))) + # weight_loss = energy_weight * (loss_dx_dy + loss_distance + loss_distance_dx_dy + loss_angular_diff) + + + + energy_weight=k*torch.exp(1+(np.log10(1/k)*(energy-e_thrs))) + # weight_loss = energy_weight * (loss_dx_dy + loss_distance + loss_distance_dx_dy + loss_angular_diff) + weight_loss = energy_weight * (loss_dx_dy + loss_distance + loss_distance_dx_dy) + + loss_distance=loss_distance.mean() + loss_dx_dy = loss_dx_dy.mean() + loss_distance_dx_dy = loss_distance_dx_dy.mean() + + loss_angular_diff = loss_angular_diff.mean() + + step_loss = step_loss + weight_loss.mean() + # step_loss = step_loss + loss_dx_dy + loss_distance + loss_distance_dx_dy + loss = loss + step_loss + + return loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_sky_direction_loss_diffusion(self, x, y, labels_energy_value): + + loss = 0 + y = y.squeeze(-1) + y_embed = self.model.target_embedder(y) + + for t in range(self.model.T): + # Add noise to target (landmarks) + alpha_bar_t = self.model.alpha_bar[t] + noise = torch.randn_like(y_embed) + z_t = torch.sqrt(alpha_bar_t) * y_embed + torch.sqrt(1 - alpha_bar_t) * noise + + # Denoise step + z, _ = self.model.blocks[t](x, z_t, None) # W_embed not needed + + preds = self.model.regress(z) + # Loss to clean target + loss_l2 = F.mse_loss(preds, y) + + # Weighted by SNR difference + step_loss = 2.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + # Final step: use regression head + if t == self.model.T - 1: + + + labels_dx_dy = y[:, 0:2] + label_distance = y[:, 2] + direction_pred = self.model.regress(z) + + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + # direction_pred = list(direction_pred) + # pred_az_atl = direction_pred[0][:,0:2] + # pred_separation = direction_pred[0][:,2] + # # direction_pred[0]= direction_pred[0][:,0:2] + # else: + + # pred_az_atl = direction_pred[:, 0:2] + # pred_separation = direction_pred[:, 2] + + # labels_az_alt = labels_direction[:, 0:2] + # label_separation = labels_direction[:, 2] + # loss_separation = self.criterion_direction(pred_separation, label_separation) + + # # loss_vector = self.criterion_vector(pred_dir_cartesian, labels_direction_cartesian) + # # vect_magnitud = torch.sqrt(torch.sum(pred_dir_cartesian**2, dim=1)) + # # loss_magnitud = torch.abs(1.0-vect_magnitud).sum() + + # # alt_az = utils_torch.cartesian_to_alt_az(direction[:,0:3]) + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + # loss_alt_az = self.criterion_alt_az(direction_pred, labels_direction) + # else: + # loss_alt_az = self.criterion_alt_az_l1(pred_az_atl, labels_az_alt) + + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + + # loss_angular_error, _ = AngularDistance( + # (direction_pred[0][:, 1]), + # labels_az_alt[:, 1], + # (direction_pred[0][:, 0]), + # labels_az_alt[:, 0], + # reduction="sum", + # ) + + # else: + + # loss_angular_error, _ = AngularDistance( + # (direction_pred[:, 1]), + # labels_az_alt[:, 1], + # (direction_pred[:, 0]), + # labels_az_alt[:, 0], + # reduction="sum", + # ) + + # if training == False: + # if len(direction_pred)>1 and type(direction_pred)!=torch.Tensor: + # _, angular_diff = AngularDistance( + # (direction_pred[0][:, 1]), + # labels_az_alt[:, 0], + # (direction_pred[0][:, 0]), + # labels_az_alt[:, 1], + # reduction=None, + # ) + # else: + # _, angular_diff = AngularDistance( + # (direction_pred[:, 1]), + # labels_az_alt[:, 0], + # (direction_pred[:, 0]), + # labels_az_alt[:, 1], + # reduction=None, + # ) + # else: + # angular_diff = None + + # loss = loss_alt_az + 0.001*(loss_separation + loss_angular_error) + + # loss = loss + step_loss + + # return loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff + # ---------------------------------------------------------------------------------------------------------- + def compute_class_weights(self, y, n_classes): + # y: tensor de etiquetas, por ejemplo torch.tensor([0,0,1,1,1,2]) + counts = torch.bincount(y, minlength=n_classes) + total = counts.sum() + # Evita división por cero: + counts = torch.clamp(counts, min=1) + weights = total / (counts * n_classes) + return weights + # ---------------------------------------------------------------------------------------------------------- + def compute_type_loss_diffusion(self, x,y, training=False): + + loss = 0 + # Get the embedding of the true label (e.g. "this is a 3") + uy = self.model.W_embed[y] + for t in range(self.model.T): + + # Get the current noise level from the schedule + alpha_bar_t = self.model.alpha_bar[t] + + # Generate random noise with the same shape as uy + noise = torch.randn_like(uy) + + # Add noise to the label embedding → this is our "noisy target" + z_t = torch.sqrt(alpha_bar_t) * uy + torch.sqrt(1 - alpha_bar_t) * noise + + # Pass the image and the noisy label through one block + z_pred, _ = self.model.blocks[t](x, z_t, self.model.W_embed) + + # Compute how far the output is from the clean label embedding + loss_l2 = F.mse_loss(z_pred, uy) + + # Weight the loss using the signal-to-noise ratio + step_loss = 0.5 * self.model.eta * self.model.snr_diff[t] * loss_l2 + + + accuracy = 0 + precision = 0 + predicted = None + # If it's the final layer, add classification and KL losses + if t == self.model.T - 1: + # Get predictions from classifier + logits = self.model.classifier(z_pred) + + # class_weights = self.compute_class_weights(y, n_classes=2).to(y.device) + # loss_ce = F.cross_entropy(logits, y, class_weights) + + # Cross-entropy loss: how wrong the predicted class is + loss_ce = F.cross_entropy(logits, y) + + # KL-like loss: penalize if embedding is too far from origin + loss_kl = 0.5 * uy.pow(2).sum(dim=1).mean() + + # Add all parts together + # loss = loss + loss_ce + loss_kl + step_loss = step_loss + loss_ce + loss_kl + classification_pred,*_ = self.model(x) + + # Calculate accuracy + predicted = torch.softmax(classification_pred, dim=1) + predicted = predicted.argmax(dim=1) + + if training: + self.class_train_accuracy.update(predicted, y) + accuracy = self.class_train_accuracy.compute().item() + self.f1_score_train.update(predicted, y) + self.precision_train.update(predicted, y) + precision = self.precision_train.compute().item() + else: + self.class_val_accuracy.update(predicted, y) + accuracy = self.class_val_accuracy.compute().item() + self.confusion_matrix.update(predicted, y) + self.f1_score_val.update(predicted, y) + self.precision_val.update(predicted, y) + precision = self.precision_val.compute().item() + loss = loss + step_loss + + return loss, accuracy, predicted, precision + # ---------------------------------------------------------------------------------------------------------- + # def on_train_start(self): + # self.dummy_tensor = self.occupy_free_gpu_memory(self.device) + # ---------------------------------------------------------------------------------------------------------- + def on_train_batch_start(self, batch, batch_idx): + if self.device != torch.device("cpu"): + if batch_idx == 5 and not hasattr(self, "dummy_tensor"): + self.dummy_tensor = self.occupy_free_gpu_memory(self.device) + + # ---------------------------------------------------------------------------------------------------------- + def training_step(self, batch, batch_idx): + # ------------------------------------------------------------------ + # Read inputs (features) and labels + # ------------------------------------------------------------------ + features, labels, t = batch + loss = 0 + if len(features) > 0: + imgs = features["image"] + + if self.task == Task.type: + labels_class = labels["type"] + + labels_energy_value = labels["energy"] + if self.task == Task.energy: + labels_energy_value = labels_energy_value.to(self.device) + + + if self.task == Task.cameradirection: + labels_direction = labels["direction"] + + imgs = imgs.to(self.device) + + # Log sample images to TensorBoard every 5 epochs + if batch_idx == 0 and self.current_epoch % 5 == 0 and self.trainer.is_global_zero: + import torchvision + try: + num_samples = min(4, imgs.shape[0]) + sample_imgs = imgs[:num_samples] + + label_str = "" + if self.task == Task.type: + label_str = str(labels_class[:num_samples].tolist()) + elif self.task == Task.energy: + label_str = str(labels_energy_value[:num_samples].tolist()) + elif self.task == Task.cameradirection: + label_str = str(labels_direction[:num_samples].tolist()) + + grid = torchvision.utils.make_grid(sample_imgs, normalize=True) + self.logger.experiment.add_image(f"Train/Images_labels_{label_str}", grid, self.current_epoch) + except Exception as e: + print(f"Error logging train images: {e}") + + # ------------------------------------------------------------------ + # Predictions based on one backbone or two back bones + # ------------------------------------------------------------------ + if not self.is_difussion: + classification_pred, energy_pred, direction_pred = self.model(imgs) + + # ------------------------------------------------------------------ + # Particle type + # --------------------------------------- + if self.task == Task.type: + if self.is_difussion: + loss, accuracy, predicted, precision = self.compute_type_loss_diffusion(imgs,labels_class,training=True) + + else: + classification_pred_ = classification_pred[0] + feature_vector = classification_pred[1] + + loss, accuracy, predicted, precision = self.compute_type_loss( + classification_pred_, labels_class, test_val=False, training=True + ) + # Log batch loss and accuracy on the progress bar + self.log( + "train_acc", + accuracy * 100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + ) + self.log( + "train_precision", + precision*100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + ) + + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection: + + if self.is_difussion: + if self.num_inputs == 2: + peak_time = features["peak_time"] + peak_time = peak_time.to(self.device) + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff = self.compute_camera_direction_loss_diffusion(imgs, labels_direction,labels_energy_value, peak_time) + else: + + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff = self.compute_camera_direction_loss_diffusion(imgs, labels_direction,labels_energy_value) + else: + if len(direction_pred)==2: + direction_pred = direction_pred[0] + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_diff = self.compute_camera_direction_loss( direction_pred, labels_direction, training=True) + + self.loss_train_distance += loss_distance.item() + self.loss_train_dx_dy += loss_dx_dy.item() + self.loss_train_angular_error += loss_angular_diff.item() + + self.loss_val_separation = 0.0 + self.loss_val_dx_dy = 0.0 + self.loss_val_angular_error = 0.0 + + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy: + + if self.is_difussion: + if self.num_inputs == 1: + loss, *_ = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=True,x_2=None) + else: + peak_time = features["peak_time"] + loss, *_ = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=True,x_2=peak_time) + + else: + if len(energy_pred)==2: + energy_pred = energy_pred[0] + + loss, *_ = self.compute_energy_loss( + energy_pred, labels_energy_value, test_val=False, training=True + ) + # --------------------------------------- + + # L2 Regularization + l2_norm = sum(p.pow(2.0).sum() for p in self.parameters()) + loss = loss + self.l2_lambda * l2_norm + + # Display on Bar progress + self.log( + "train_loss", + loss.item(), + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + ) + + if not np.isnan(loss.item()): + + self.loss_train_sum += loss.item() + self.num_train_batches += 1 + + # torch.cuda.empty_cache() + return loss + # ---------------------------------------------------------------------------------------------------------- + def on_train_epoch_end(self): + + if self.trainer.world_size > 1: + + # dist.barrier() + + # Gather y Sync values with the GPUs. + total_loss_train= self.all_gather(self.loss_train_sum).sum().item() + total_batches_val = self.all_gather(torch.tensor(self.num_train_batches, device=self.device)).sum().item() + + else: + total_loss_train = self.loss_train_sum + total_batches_val = self.num_train_batches + + if total_batches_val==0: + + return 0 + + if self.trainer.is_global_zero: + + global_loss = total_loss_train / total_batches_val + + self.logger.experiment.add_scalars( + "loss/Global Training Loss", + { + "loss": global_loss, + }, + self.current_epoch, + ) + + self.logger.experiment.add_scalars( + "Learning rate", + { + "Learning rate": self.scheduler.get_last_lr()[0], + }, + self.current_epoch, + ) + # Optionally print the global loss for immediate feedback + print(f"Epoch {self.current_epoch}: Global Training Loss: {global_loss:.4f}") + + # --------------------------------------- + # Particle Type + # --------------------------------------- + if self.task == Task.type: + ii = 0 + # f1_score = self.f1_score_train.compute().detach().cpu().numpy()*100.0 + # self.f1_score_train.reset() + + # precision = self.precision_train.compute().detach().cpu().numpy()*100.0 + # self.precision_train.reset() + + epoch_accuracy = self.class_train_accuracy.compute().detach().cpu().item() * 100 + self.class_train_accuracy.reset() + + # # Compute the accuracy and reset the metric states after each epoch + if self.trainer.is_global_zero: + # self.log("train_acc_epoch", epoch_accuracy, on_step=False, prog_bar=True, sync_dist=True) + # # Log + # self.logger.experiment.add_scalars( + # "Metrics/Training", + # { + # "acc": epoch_accuracy, + # "f1":f1_score, + # "precision":precision, + # }, + # self.current_epoch, + # ) + print( + f"Epoch {self.current_epoch}: Global Training Accuracy: {epoch_accuracy:.4f}" + ) + filename_prefix = ( + f"Epoch_{self.current_epoch}_{self.task.name}_train_acc" + ) + if self.logger.log_dir: + self.save_checkpoint( + self.logger.log_dir, + epoch_accuracy, + filename_prefix=filename_prefix, + is_loss=False, + ) + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection and self.trainer.is_global_zero: + + filename_prefix = ( + f"Epoch_{self.current_epoch}_{self.task.name}_train_loss" + ) + self.save_checkpoint( + self.logger.log_dir, + global_loss, + filename_prefix=filename_prefix, + is_loss=True, + ) + # Log scalar values + self.logger.experiment.add_scalars( + "loss/ Loss Training", + { + "loss": global_loss, + "loss_distance": self.loss_train_distance + / self.num_train_batches, + "loss_alt_az": self.loss_train_dx_dy / self.num_train_batches, + "loss_angular_error": self.loss_train_angular_error + / self.num_train_batches, + }, + self.current_epoch, + ) + self.loss_train_distance = 0.0 + self.loss_train_dx_dy = 0.0 + self.loss_train_angular_error = 0.0 + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy and self.trainer.is_global_zero: + filename_prefix = ( + f"Epoch_{self.current_epoch}_{self.task.name}_train_loss" + ) + self.save_checkpoint( + self.logger.log_dir, + global_loss, + filename_prefix=filename_prefix, + is_loss=True, + ) + # --------------------------------------- + # Delete all the lists and set to 0 + # the values used to estimate the losses + # --------------------------------------- + self.reset_values() + print("END train epoch") + # Reset + self.loss_train_sum = 0 + self.num_train_batches = 0 + self.training_step_outputs = [] + # ---------------------------------------------------------------------------------------------------------- + @torch.no_grad() + def validation_step(self, batch, batch_idx, dataloader_idx=0): + loss=0 + self.model.eval() + + # ------------------------------------------------------------------ + # Read inputs (features) and labels + # ------------------------------------------------------------------ + features, labels, t = batch + + if len(features) > 0: + + imgs = features["image"] + + batch_size = imgs.shape[0] + if self.task == Task.type: + labels_class = labels["type"] + + labels_energy_value = labels["energy"] + hillas_intensity = features["hillas"]["hillas_intensity"] + + if self.task == Task.cameradirection: + labels_direction = labels["direction"] + + + # Log sample images to TensorBoard every 5 epochs + if batch_idx == 0 and self.current_epoch % 5 == 0 and self.trainer.is_global_zero: + import torchvision + try: + num_samples = min(4, imgs.shape[0]) + sample_imgs = imgs[:num_samples].to(self.device) + + label_str = "" + if self.task == Task.type: + label_str = str(labels_class[:num_samples].tolist()) + elif self.task == Task.energy: + label_str = str(labels_energy_value[:num_samples].tolist()) + elif self.task == Task.cameradirection: + label_str = str(labels_direction[:num_samples].tolist()) + + grid = torchvision.utils.make_grid(sample_imgs, normalize=True) + self.logger.experiment.add_image(f"Validation/Images_labels_{label_str}", grid, self.current_epoch) + except Exception as e: + print(f"Error logging validation images: {e}") + # ------------------------------------------------------------------ + # Predictions based on one backbone or two back bones + # ------------------------------------------------------------------ + if not self.is_difussion: + classification_pred, energy_pred, direction_pred = self.model(imgs) + # ------------------------------------------------------------------ + # Compute Loss functions based on different tasks + # ------------------------------------------------------------------ + # Particle Type + # --------------------------------------- + if self.task == Task.type: + if self.is_difussion: + loss, accuracy, predicted, precision = self.compute_type_loss_diffusion(imgs,labels_class,training=False) + else: + classification_pred_ = classification_pred[0] + feature_vector = classification_pred[1] + loss, accuracy, predicted, precision = self.compute_type_loss( + classification_pred_, + labels_class, + test_val=False, + training=False, + ) + self.val_proton_gammanes.extend( + torch.softmax(classification_pred_, dim=1)[:, 1][labels_class == 0] + .float() + .cpu() + .detach() + .numpy() + .flatten() + .tolist() + ) + self.val_gamma_gammanes.extend( + torch.softmax(classification_pred_, dim=1)[:, 1][labels_class == 1] + .float() + .cpu() + .detach() + .numpy() + .flatten() + .tolist() + ) + # Log batch loss and accuracy on the progress bar + # if dataloader_idx == 0: + # if self.is_difussion: + # loss, accuracy, predicted, precision =self.compute_type_loss_diffusion( + # classification_pred_, + # labels_class, + # test_val=False, + # training=False, + # ) + # else: + # loss, accuracy, predicted, precision = self.compute_type_loss( + # classification_pred_, + # labels_class, + # test_val=False, + # training=False, + # ) + + self.log( + "val_acc", + accuracy * 100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + batch_size=batch_size, + ) + self.log( + "val_prec", + precision*100, + on_step=True, + on_epoch=False, + prog_bar=True, + logger=False, + batch_size=batch_size, + ) + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection: + + if len(direction_pred)==2: + direction_pred = direction_pred[0] + + loss, loss_dx_dy, loss_distance, loss_distance_dx_dy, loss_angular_diff, angular_error= self.compute_camera_direction_loss( direction_pred, labels_direction, training=False) + + # ------------------------------------------------------------------------ + # Convert the offset to altitud and azimuth + # ------------------------------------------------------------------------ + + if dataloader_idx == 0: + self.loss_val_distance += loss_distance.item() + self.loss_val_dx_dy += loss_dx_dy.item() + self.loss_val_angular_error += loss_angular_diff.item() + self.val_angular_diff_list.extend(angular_error) + # ------------------------------------------------------- + + pred_dx = direction_pred[:, 0].float().cpu().detach().numpy() + pred_dy = direction_pred[:, 1].float().cpu().detach().numpy() + + cam_x = pred_dx + cam_y = pred_dy + + pred_alt, pred_az = self.val_dataset.cam_to_alt_az(labels["tel_ids"].cpu().detach().numpy(), labels["focal_length"].cpu().detach().numpy(), labels["pix_rotation"].cpu().detach().numpy(),labels["tel_az"].cpu().detach().numpy(),labels["tel_alt"].cpu().detach().numpy(), cam_x, cam_y) + + labels_dx_dy = labels_direction[:, 0:2] + + true_alt, true_az = self.val_dataset.cam_to_alt_az(labels["tel_ids"].cpu().detach().numpy(), labels["focal_length"].cpu().detach().numpy(), labels["pix_rotation"].cpu().detach().numpy(),labels["tel_az"].cpu().detach().numpy(),labels["tel_alt"].cpu().detach().numpy(), labels_dx_dy[:,0].float().cpu().detach().numpy(), labels_dx_dy[:,1].float().cpu().detach().numpy()) + + self.val_alt_pred_list.extend(np.radians(pred_alt)) + self.val_az_pred_list.extend(np.radians(pred_az)) + self.val_alt_label_list.extend(np.radians(true_alt)) + self.val_az_label_list.extend(np.radians(true_az)) + + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy: + if self.is_difussion: + if self.num_inputs == 1: + loss, energy_diff, energy_pred_tev = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=False,x_2=None) + else: + peak_time = features["peak_time"] + loss, energy_diff, energy_pred_tev = self.compute_energy_loss_diffusion(imgs,labels_energy_value/1.0,training=False,x_2=peak_time) + else: + if len(energy_pred)==2: + energy_pred = energy_pred[0] + + loss, energy_diff = self.compute_energy_loss( + energy_pred, labels_energy_value/1.0, test_val=False, training=False + ) + energy_pred_tev = torch.pow(10, energy_pred*1.0) + + if dataloader_idx == 0: + + self.val_energy_diff_list.extend(energy_diff) + self.val_energy_pred_list.extend( + energy_pred_tev[:, 0].float().cpu().detach().numpy() + ) + + # --------------------------------------- + # Collect the True Energy and Hillas Intensity + # --------------------------------------- + + energy_label_tev = torch.pow(10, labels_energy_value) + + if dataloader_idx == 0: + self.val_energy_label_list.extend( + energy_label_tev[:, 0].float().cpu().detach().numpy().flatten().tolist() + ) + self.val_hillas_intensity_list.extend( + hillas_intensity[:, 0].float().cpu().detach().numpy().flatten().tolist() + ) + + # --------------------------------------- + # Log validation loss + # --------------------------------------- + + if dataloader_idx == 0: + self.loss_val_sum += loss.item() + self.num_val_batches += 1 + + + loss_key = "val_loss" if dataloader_idx == 0 else "test_loss" + if loss is not None: + self.log(loss_key, loss,on_step=False, on_epoch=True, prog_bar=True, logger=True, sync_dist=True, batch_size=batch_size) + + # Release cuda memory + # torch.cuda.empty_cache() + return loss + # ---------------------------------------------------------------------------------------------------------- + @torch.no_grad() + def on_validation_epoch_end(self): + + if self.trainer.world_size > 1: + dist.barrier() + # Gather y Sync values with the GPUs. + total_loss_val = self.all_gather(self.loss_val_sum).sum().item() + total_batches_val = self.all_gather(torch.tensor(self.num_val_batches, device=self.device)).sum().item() + + else: + total_loss_val = self.loss_val_sum + total_batches_val = self.num_val_batches + + + if self.trainer.is_global_zero: + # Calcular la pérdida promedio global + global_loss_val = total_loss_val / max(1, total_batches_val) + + self.logger.experiment.add_scalars( + "loss/Global Validation Loss", + { + "loss": global_loss_val, + }, + self.current_epoch, + ) + + # Print the global loss for immediate feedback + print(f"Epoch {self.current_epoch}: Global Validation Loss: {global_loss_val:.4f}") + # --------------------------------------- + # Particle Type + # --------------------------------------- + if self.task == Task.type: + conf_matrix = self.confusion_matrix.compute().detach().cpu().numpy() + f1_score_val = self.f1_score_val.compute().detach().cpu().numpy()*100.0 + precision_val = self.precision_val.compute().detach().cpu().numpy()*100.0 + + # Compute the accuracy and reset the metric states after each epoch + epoch_accuracy_val = self.class_val_accuracy.compute().item() * 100 + + self.class_val_accuracy.reset() + self.confusion_matrix.reset() + self.f1_score_val.reset() + self.precision_val.reset() + + # --------------------------------------- + # Create Confusion Matrix + # --------------------------------------- + if self.trainer.is_global_zero: + + # Log + self.logger.experiment.add_scalars( + "Metrics/Validation", + { + "acc": epoch_accuracy_val, + "f1": f1_score_val, + "precision":precision_val, + }, + self.current_epoch, + ) + + print( + f"Epoch {self.current_epoch}: Global Validation Accuracy: {epoch_accuracy_val:.4f}" + ) + print( + f"Epoch {self.current_epoch}: Global Validation F1 Score: {f1_score_val:.4f}" + ) + print( + f"Epoch {self.current_epoch}: Global Validation Precision: {precision_val:.4f}" + ) + filename_prefix = "confusion_matrix_val" + cm_file_name = f"{filename_prefix}_{self.current_epoch}_{epoch_accuracy_val:.4f}_Validation" + + if self.logger.log_dir: + plot_confusion_matrix( + conf_matrix, + self.class_names, + cm_file_name, + self.logger.log_dir, + ) + # Compute class-wise accuracies + class_accuracies = (conf_matrix.diagonal() / conf_matrix.sum(axis=1))*100 + + self.logger.experiment.add_scalars( + "confusion_matrix", + { + "val_acc_gamma": class_accuracies[0], + "val_acc_proton": class_accuracies[1], + "val_global_acc": epoch_accuracy_val, + }, + self.current_epoch, + ) + # Add the protonness and gammaness histogramns + self.val_proton_gammanes = np.array(self.val_proton_gammanes) + self.val_gamma_gammanes = np.array(self.val_gamma_gammanes) + if self.trainer.is_global_zero: + fig, ax = plt.subplots(figsize=(8, 6)) + ax.hist( + self.val_proton_gammanes, + bins=50, + alpha=0.7, + label="Proton", + color="blue", + density=True, + ) + ax.hist( + self.val_gamma_gammanes, + bins=50, + alpha=0.7, + label="Gammas", + color="orange", + density=True, + ) + ax.set_xlabel("Score") + ax.set_ylabel("Density") + ax.set_title("Protons and Gammas Gammanes Distribution - Validation") + ax.legend() + save_path = os.path.join( + self.logger.log_dir, + f"proton_gammaness_validation_{self.current_epoch}_{global_loss_val}.png" + ) + fig.savefig(save_path, format="png") # <--- Usa fig.savefig, no plt.savefig + + # 2. Luego envíala al logger + self.logger.experiment.add_figure( + "Protons and Gammas Gammanes Distribution - Validation", + fig, + self.current_epoch, + ) + + # 3. Finalmente cierra la figura + plt.close(fig) + + # Clear lists + self.val_proton_gammanes = [] + self.val_gamma_gammanes = [] + # return 0 + # --------------------------------------- + # Direction + # --------------------------------------- + if self.task == Task.cameradirection and self.trainer.is_global_zero: + + self.print_direction_error(self.val_angular_diff_list, "Validation") + + plt.close("all") + fig_direction_error = plot_direction_resolution_error( + self.val_alt_pred_list, + self.val_az_pred_list, + self.val_alt_label_list, + self.val_az_label_list, + self.val_energy_label_list, + self.val_hillas_intensity_list, + ) + self.logger.experiment.add_figure( + "Direction Resolution Error/Validation", + fig_direction_error, + self.current_epoch, + ) # Log the plot + + # Save the figure + fig_direction_error.savefig( + os.path.join( + self.logger.log_dir, + "angular_resolution_validation_" + + str(self.current_epoch) + + "_" + + str(global_loss_val) + + ".png", + ), + format="png", + ) + plt.close(fig_direction_error) # Close the figure to release memory + plt.close("all") + + # Log scalar values + self.logger.experiment.add_scalars( + "loss/Loss Validation", + { + "loss": global_loss_val, + "loss_distance": self.loss_val_distance / self.num_val_batches, + "loss_dx_dy": self.loss_val_dx_dy / self.num_val_batches, + "loss_angular_error": self.loss_val_angular_error + / self.num_val_batches, + }, + self.current_epoch, + ) + + + + # --------------------------------------- + # Energy + # --------------------------------------- + if self.task == Task.energy and self.trainer.is_global_zero: + + self.print_energy_error(self.val_energy_diff_list, "Validation") + plt.close("all") + fig_energy_error = plot_energy_resolution_error( + self.val_energy_pred_list, + self.val_energy_label_list, + self.val_hillas_intensity_list, + ) + self.logger.experiment.add_figure( + "Energy Resolution Error/Validation", + fig_energy_error, + self.current_epoch, + ) # Log the plot + fig_energy_error.savefig( + os.path.join( + self.logger.log_dir, + "error_resolution_validation_" + + str(self.current_epoch) + + "_" + + str(global_loss_val) + + ".png", + ), + format="png", + ) + + plt.close(fig_energy_error) # Close the figure to release memory + plt.close("all") + + # Add histograms for energy labels and predictions + fig_hist, ax_hist = plt.subplots(figsize=(8, 6)) + + import numpy as np + # Use percentiles of true labels to avoid extreme predictions ruining the x-axis range + min_val = np.percentile(self.val_energy_label_list, 1) + max_val = np.percentile(self.val_energy_label_list, 99) + + # Energy typically spans multiple orders of magnitude, so log bins are more appropriate + min_val = max(1e-3, min_val) + if min_val >= max_val: + max_val = min_val * 10 + shared_bins = np.logspace(np.log10(min_val), np.log10(max_val), 50) + + ax_hist.hist( + self.val_energy_label_list, + bins=shared_bins, + alpha=0.5, + label="Labels", + color="blue", + ) + ax_hist.hist( + self.val_energy_pred_list, + bins=shared_bins, + alpha=0.5, + label="Predictions", + color="orange", + ) + ax_hist.set_xscale("log") + ax_hist.set_xlabel("Energy (TeV)") + ax_hist.set_ylabel("Counts") + ax_hist.set_title("Energy Distribution - Labels vs Predictions (Validation)") + ax_hist.legend() + + self.logger.experiment.add_figure( + "Energy Distribution - Labels vs Predictions/Validation", + fig_hist, + self.current_epoch, + ) + + save_path = os.path.join( + self.logger.log_dir, + f"energy_distribution_validation_{self.current_epoch}_{global_loss_val}.png" + ) + fig_hist.savefig(save_path, format="png") + plt.close(fig_hist) + plt.close("all") + + # ---------------------------------------------------------------------------------------------------------- + def reset_values(self): + # --------------------------------------- + # Reset + # --------------------------------------- + self.loss_val_sum = 0 + self.num_val_batches = 0 + + # Reset Energy and hillas intensity + # -------------------------------------------------- + # Validation + # -------------------------------------------------- + self.val_energy_label_list.clear() + self.val_hillas_intensity_list.clear() + + # -------------------------------------------------- + # Reset Direction + # -------------------------------------------------- + # Validation + # -------------------------------------------------- + self.val_angular_diff_list.clear() + + + self.val_alt_pred_list.clear() + self.val_az_pred_list.clear() + self.val_alt_label_list.clear() + self.val_az_label_list.clear() + + self.loss_val_distance = 0.0 + self.loss_val_dx_dy = 0.0 + self.loss_val_angular_error = 0.0 + + # -------------------------------------------------- + # Reset Energy + # -------------------------------------------------- + # Validation + # -------------------------------------------------- + self.val_energy_diff_list.clear() + self.val_energy_pred_list.clear() + + + # ---------------------------------------------------------------------------------------------------------- + def print_direction_error(self, angular_diff_list, type_val: str): + # Count the angular error in ranges [20º-0.1º] + error_20 = len([num for num in angular_diff_list if num < 20]) + error_10 = len([num for num in angular_diff_list if num < 10]) + error_5 = len([num for num in angular_diff_list if num < 5]) + error_2 = len([num for num in angular_diff_list if num < 2]) + error_1 = len([num for num in angular_diff_list if num < 1]) + error_0_5 = len([num for num in angular_diff_list if num < 0.5]) + error_0_25 = len([num for num in angular_diff_list if num < 0.25]) + error_0_1 = len([num for num in angular_diff_list if num < 0.1]) + # Log + self.logger.experiment.add_scalars( + "Direction Error/" + type_val, + { + "0: error: 20": error_20, + "1: error: 10": error_10, + "2: error: 5": error_5, + "3: error: 2": error_2, + "4: error: 1": error_1, + "5: error: 0.5": error_0_5, + "6: error: 0.25": error_0_25, + "7: error: 0.1": error_0_1, + }, + self.current_epoch, + ) + # Print + print(type_val + " Direction Error < 20º: ", error_20) + print(type_val + " Direction Error < 10º: ", error_10) + print(type_val + " Direction Error < 5º: ", error_5) + print(type_val + " Direction Error < 2º: ", error_2) + print(type_val + " Direction Error < 1º: ", error_1) + print(type_val + " Direction Error < 0.50º: ", error_0_5) + print(type_val + " Direction Error < 0.25º: ", error_0_25) + print(type_val + " Direction Error < 0.10º: ", error_0_1) + # ---------------------------------------------------------------------------------------------------------- + def print_energy_error(self, energy_diff_list, type_val: str): + + + error_30 = len([num for num in energy_diff_list if num < 30]) + error_20 = len([num for num in energy_diff_list if num < 20]) + error_10 = len([num for num in energy_diff_list if num < 10]) + error_5 = len([num for num in energy_diff_list if num < 5]) + error_1 = len([num for num in energy_diff_list if num < 1]) + error_0_5 = len([num for num in energy_diff_list if num < 0.5]) + error_0_25 = len([num for num in energy_diff_list if num < 0.25]) + error_0_1 = len([num for num in energy_diff_list if num < 0.1]) + error_0_01 = len([num for num in energy_diff_list if num < 0.01]) + # Log + self.logger.experiment.add_scalars( + "Energy Error/" + type_val, + { + "0: error: 30": error_30, + "1: error: 20": error_20, + "2: error: 10": error_10, + "3: error: 5": error_5, + "4: error: 1": error_1, + "5: error: 0.5": error_0_5, + "6: error: 0.25": error_0_25, + "7: error: 0.1": error_0_1, + "8: error: 0.01": error_0_01, + }, + self.current_epoch, + ) + # Print + print(f"Total: {len(energy_diff_list)}") + print(type_val + " Energy Error < 30:", error_30) + print(type_val + " Energy Error < 20:", error_20) + print(type_val + " Energy Error < 10:", error_10) + print(type_val + " Energy Error < 5:", error_5) + print(type_val + " Energy Error < 1:", error_1) + print(type_val + " Energy Error < 0.5:", error_0_5) + print(type_val + " Energy Error < 0.1:", error_0_1) + print(type_val + " Energy Error < 0.01:", error_0_01) + # ---------------------------------------------------------------------------------------------------------- + def create_confusion_matrix( + self, epoch, accuracy, all_val_preds, all_val_labels, val_type: str + ): + all_labels = all_val_preds.numpy() + all_preds = all_val_labels.numpy() + filename_prefix = "confusion_matrix_val" + cm_file_name = f"{filename_prefix}_{epoch}_{accuracy:.4f}_{val_type}" + cm = confusion_matrix(all_labels, all_preds) + + cm_norm = cm.astype("float") / cm.sum(axis=1)[:, np.newaxis] + + # Remove NaNs + cm_norm = np.nan_to_num(cm_norm, nan=0.0) + accuracies = cm_norm.diagonal() * 100.0 + if self.logger.log_dir: + plot_confusion_matrix( + cm, + self.class_names, + accuracies, + cm_file_name, + self.logger.log_dir, # self.save_folder + ) + + return accuracies + # ---------------------------------------------------------------------------------------------------------- + def configure_optimizers(self): + + if self.optimizer_type == "sgd": + print("Using SGD...") + # Optimizer SGD + self.optimizer = optim.SGD( + self.model.parameters(), + lr=self.learning_rate, + momentum=self.momentum, + weight_decay=self.weight_decay, + nesterov=True, + ) + elif self.optimizer_type == "adam": + print("Using Adam...") + # Optimizer Adam + self.optimizer = optim.Adam( + self.model.parameters(), + lr=self.learning_rate, + eps=self.adam_epsilon, + weight_decay=self.weight_decay, + ) + elif self.optimizer_type == "adamw": + print("Using AdamW...") + # Optimizer Adam + self.optimizer = optim.AdamW( + self.model.parameters(), + lr=self.learning_rate, + eps=self.adam_epsilon, + weight_decay=self.weight_decay, + ) + else: + raise ValueError( + f"Unsupported optimizer type: {self.optimizer_type}. Supported types are 'sgd' and 'adam'." + ) + + # Cosine phase + # Note: The schedule is updated after execute one epoch. -> self.epochs* self.batches + lf = one_cycle( + 1, self.lrf, self.num_epochs * self.steps_epoch + ) # cosine 1->hyp['lrf'] + # self.scheduler = lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lf) + + # TODO: Add the option to choose the kind of scheduler + # self.scheduler = torch.optim.lr_scheduler.OneCycleLR(self.optimizer,max_lr=self.lrf,total_steps=self.num_epochs * self.steps_epoch) + self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + self.optimizer, T_max=self.num_epochs + ) + self.scheduler.last_epoch = self.start_epoch + # Prepare scheduler dictionary as expected by PyTorch Lightning + lr_dict = { + "scheduler": self.scheduler, + "interval": "epoch", # Use for each epoch, dont use for each step + "frequency": 1, + "name": "OneCycleLR", + "monitor": "val_loss", + } + # Return the optimizer and learning rate scheduler + return { + "optimizer": self.optimizer, + "lr_scheduler": lr_dict, + } + # ---------------------------------------------------------------------------------------------------------- diff --git a/ctlearn/tools/train/pytorch/__init__.py b/ctlearn/tools/train/pytorch/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/ctlearn/tools/train/pytorch/config/default_config_file.yml b/ctlearn/tools/train/pytorch/config/default_config_file.yml new file mode 100644 index 00000000..abf55871 --- /dev/null +++ b/ctlearn/tools/train/pytorch/config/default_config_file.yml @@ -0,0 +1,152 @@ +data: + train_gamma_proton: ./data/gamma_proton_train_remix.dl1.pickle + validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle + + train_gamma: ./data/gamma_955000_train.pickle + validation_gamma: ./data/gamma_106141_validation.pickle + + test_gamma: ./data/gamma_1805522_test_gamma.pickle + test_proton: ./data/proton_130811_test_proton.pickle + test_electron: None + test_validation_gamma: ./data/gamma_180552_test_val_gamma.pickle + + test_validation_gamma_proton: ./data/gamma_proton_212282_validation.pickle + + observation: ./run_2931.dl1.pickle + # Important: This is only for testing purpose. Set always to 0 + # when you are training, validating or estimating the dl2 files + training_reduce_factor: 0 #64 #4 + validation_reduce_factor: 0 #16 #8 + validation_test_reduce_factor: 0 #16 #8 + + # Check points + type_checkpoint: ./run/run_type_training_14/exp_14_type_train/version_0/Epoch_6_type_train_acc_80.9682309627532959.pth + energy_checkpoint: /home/cpozogonzalez/ctlearn/run/run_energy_training_14/exp_14_energy_train/version_1/Epoch_13_energy_train_loss_30.2185532478501244.pth + direction_checkpoint: /lhome/ext/ucm147/ucm1477/data/check_points/v_5/Epoch_23_cameradirection_train_loss_7296.6534562211982120.pth + +run_details: + + mode: "observation" # The option are: "train", "results", "observation" and "validate" + task: "direction" # The option are: "all", "energy" "type" and "direction" + test_type: "gamma" # The option are: "gamma" "proton" or "electron" + experiment_number: 14 # The experiment number. The experiment folder is saved into the "run" folder. + + +cut-off: + + leakage_intensity: 0.2 # bigger to this value, the event is removed + intensity: 50 # below to this value, the event is removed + +model: + + model_type: + model_name: "DoubleBBEfficientNet" + parameters: + model_variant: "efficientnet-b3" + task: 'type' + num_outputs: 2 + device_str: "cuda" + energy_bins: None + + model_energy: + model_name: "ThinResNet" + parameters: + task: 'energy' + num_inputs: 1 + num_outputs: 1 + num_blocks: [3, 4, 6, 3] #[2, 3, 3, 3] + dropout: 0.1 + use_bn: False + + model_direction: + model_name: "ThinResNet_DBB" + parameters: + task: 'direction' + num_inputs: 1 + num_outputs: 3 + num_blocks: [3, 4, 6, 3] + dropout: 0.1 + use_bn: False + + # model_direction: + # model_name: "DBBNoPropDTReg" + # parameters: + # task: 'direction' + # num_outputs: 3 + # embedding_dim: 512 + # T: 3 + # eta: 0.1 #0.1 + +# Hyper-parameters +hyp: + + epochs: 30 + batches: 128 #128 #64 + dynamic_batches: True + optimizer: Adamw + momentum: 0.957 #Yolo 0.937 # Efficient-b3 0.757 + weight_decay: 0.0005 #0.004676 #0.0001 #0.00002 Efficient-b3 0.0005 + learning_rate: 1e-4 #1e-5 #Efficient-b3 1e-5 + lrf: 0.1 + start_epoch: 0 + steps_epoch: 100 # Computed online. Must be removed + l2_lambda: 1e-7 #1e-5 #1e-5 # L2 regularization (Set to 0.0 to skip the L2 Regularization) + adam_epsilon: 1.0e-08 #7.511309034256153e-05 #1.0e-08 + gradient_clip_val: 3.0 # Avoid gradient explosion + + save_k: 200 # Save as maximum k checkpoints. + +augmentation: + # probabilities for augmentation range = [0, 1.0] + # prob = 0.0 -> Always apply the augmentation + # prob >= 1.0 -> Never apply the augmentation, i.e., Set bigger than 1.0 ( ex: 2.0) if you want disable it. + # Note: mask augmentation is always on even with flag use_augmentation = True + # To disable it, just set to 2.5 for example. + + use_augmentation: True # This apply only on training mode. + aug_prob: 0.5 # Probability of use Augmentation + rot_prob: 0.5 # Rotation probability + trans_prob: 0.5 # Translation probability + flip_hor_prob: 0.5 # Horizontal Flip probability + flip_ver_prob: 0.5 # Vertical Flip probability + mask_prob: 0.5 # Apply mask probability + mask_dvr_prob: 0.5 # Apply dvr mask probability + noise_prob: 0.5 # No implemented yet. + max_rot: 5 # Maximum rotation in augmentation + max_trans: 10 # Maximum translation in augmentation + +normalization: + + # Normalization: Im' = (Im-mu)/sigma + use_clean: True # Use the image with the applied mask (True), IOC the mask is not applied (False) + use_clean_dvr: False + type_mu: 0.0 + type_sigma: 1000.0 + + dir_mu: 0.0 + dir_sigma: 1000.0 + + energy_mu: 0.0 + energy_sigma: 1000.0 + +dataset: + num_workers: 1 # + pin_memory: True + persistent_workers: True # + +# Hardware Architecture and precision +arch: + # device: 'mps' # Apple Mx + device: 'cuda' + precision_type: "32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" + precision_energy: "32-true" #"32-true" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" + precision_direction: "32-true" # "bf16-mixed" # Options: "64-true" "32-true" "16-true" "16-mixed" "bf16-mixed" "bf16-true" + # (bf16 for GPU with Ampere or higher, it is better that 16 because is numerical more stability) + # devices: [0,1] # [0,1] For multiple GPUs + devices: [0,1] + # Note: Check the documentation for more information. + strategy: 'deepspeed_stage_2' # Options: auto, dpp, dpp_swap, fsdp, deepspeed, horovod, bagua, deepspeed_stage_2, deepspeed_stage_3, colossalai, hivemind, etc... + +Notes: + Note_1: Training with augmentation dvr using 1-3 dilatations + Note_2: Trainining b3 applying always the mask \ No newline at end of file diff --git a/ctlearn/tools/train/pytorch/train_pytorch_model.py b/ctlearn/tools/train/pytorch/train_pytorch_model.py new file mode 100644 index 00000000..c818ca6d --- /dev/null +++ b/ctlearn/tools/train/pytorch/train_pytorch_model.py @@ -0,0 +1,667 @@ +from ctapipe.core.traits import Path, Bool, Unicode, ComponentName +from torch.utils.data import DataLoader +from ctlearn.tools.train.pytorch.CTLearnPL import CTLearnTrainer, CTLearnPL +from ctlearn.core.model import CTLearnModel +import ctlearn.core.pytorch.model +import sys +import os + +is_debug = '--debug' in sys.argv or any(arg.startswith('--log-level=DEBUG') for arg in sys.argv) + +if not is_debug: + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Suppress TensorFlow logging + os.environ['NCCL_DEBUG'] = 'WARN' # Suppress verbose NCCL networking logs + +import warnings +warnings.filterwarnings("ignore", ".*AccumulateGrad node's stream does not match.*") +warnings.filterwarnings("ignore", ".*torch.distributed.nn.functional.all_gather is deprecated.*") +warnings.filterwarnings("ignore", ".*does not have many workers which may be a bottleneck.*") +warnings.filterwarnings("ignore", ".*isinstance.treespec, LeafSpec.*") +warnings.filterwarnings("ignore", ".*This axis already has a converter set and is updating.*") + +if not is_debug: + # Silence noisy external library warnings + warnings.filterwarnings("ignore", ".*NoneDefaultNotAllowedWarning.*") + warnings.filterwarnings("ignore", ".*MergeConflictWarning.*") + warnings.filterwarnings("ignore", ".*'ctlearn.tools.train_model' found in sys.modules.*") + +try: + import torch + if hasattr(torch.autograd.graph, "set_warn_on_accumulate_grad_stream_mismatch"): + torch.autograd.graph.set_warn_on_accumulate_grad_stream_mismatch(False) +except ImportError: + raise ImportError("pytorch is not installed in your environment!") + +try: + from pytorch_lightning.loggers import TensorBoardLogger +except ImportError: + raise ImportError("pytorch_lightning is not installed in your environment!") + +from ctlearn.tools.train.base_train_model import TrainCTLearnModel +from ctlearn.core.ctlearn_enum import Task, Mode +from .utils import ( + str_list_to_enum_list, + sanity_check, + read_configuration, + create_experiment_folder, + expected_structure, +) + +from ctlearn.core.pytorch.net_utils import create_model, ModelHelper +from ctlearn.core.data_loader.loader import DLDataLoader +from pytorch_lightning.callbacks import Callback +import os +import numpy as np +import json + + +class GPUStatsLogger(Callback): + def on_train_epoch_end(self, trainer, pl_module): + mem_allocated = torch.cuda.memory_allocated() + mem_reserved = torch.cuda.memory_reserved() + + trainer.logger.experiment.add_scalar( + "gpu_mem_allocated", mem_allocated, global_step=trainer.current_epoch + ) + trainer.logger.experiment.add_scalar( + "gpu_mem_reserved", mem_reserved, global_step=trainer.current_epoch + ) + +# from ctlearn.tools.train_model import +class TrainPyTorchModel(TrainCTLearnModel): + """ + Tool to train a ``~ctlearn.core.model.CTLearnModel`` on R1/DL1a data using PyTorch. + + The tool sets up the PyTorch model using ... The PyTorch model is trained + on the input data (R1 calibrated waveforms or DL1a images) and saved in the output directory. + """ + + name = "ctlearn-train-pytorch-model" + description = __doc__ + + examples = """ + To train a CTLearn PyTorch model for the classification of the primary particle type: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --background /path/to/your/protons_dl1_dir/ \\ + --pattern-background "proton_*_run1.dl1.h5" \\ + --pattern-background "proton_*_run10.dl1.h5" \\ + --output /path/to/your/type/ \\ + --reco type \\ + + To train a CTLearn PyTorch model for the regression of the primary particle energy: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/energy/ \\ + --reco energy \\ + + To train a CTLearn PyTorch model for the regression of the primary particle + arrival direction based on the offsets in camera coordinates: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/direction/ \\ + --reco cameradirection \\ + + To train a CTLearn PyTorch model for the regression of the primary particle + arrival direction based on the offsets in sky coordinates: + > ctlearn-train-pytorch-model \\ + --signal /path/to/your/gammas_dl1_dir/ \\ + --pattern-signal "gamma_*_run1.dl1.h5" \\ + --pattern-signal "gamma_*_run10.dl1.h5" \\ + --output /path/to/your/direction/ \\ + --reco skydirection \\ + """ + + config_file = Path( + exits=True, + default_value=None, + allow_none=True, + directory_ok=True, + file_ok=True, + help="Configuration file.", + ).tag(config=True) + + disable_progress_bar = Bool( + default_value=False, + help="Disable PyTorch Lightning progress bar.", + ).tag(config=True) + + model_type = ComponentName( + CTLearnModel, default_value="PyTorchResNet" + ).tag(config=True) + + aliases = { + **TrainCTLearnModel.aliases, + "config_file": "TrainPyTorchModel.config_file", + "disable_progress_bar": "TrainPyTorchModel.disable_progress_bar", + "model-type": "TrainPyTorchModel.model_type", + "model-name": "TrainPyTorchModel.model_type", + } + + def __init__(self, **kwargs): + + # Setup GPU + os.environ["NCCL_P2P_DISABLE"] = "1" + os.environ["NCCL_IB_DISABLE"] = "1" + os.environ["NCCL_DEBUG"] = "WARN" + os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1" + torch.set_float32_matmul_precision('medium') + + super().__init__(**kwargs) + + + def setup(self): + + super().setup() + + # Create tasks Enum List + self.tasks = str_list_to_enum_list(self.reco_tasks) + + for task_ in self.tasks: + print("Task:", task_.name) + + if self.config_file is not None: + self.log.info("Loading configuration from legacy PyTorch config file: %s", self.config_file) + legacy_params = read_configuration(self.config_file) + sanity_check(legacy_params, expected_structure) + + def get_conf_val(key1, key2, trait_name, default_val): + in_config = False + for cls_name in ["TrainPyTorchModel", "TrainCTLearnModel"]: + if cls_name in self.config and trait_name in self.config[cls_name]: + in_config = True + break + if not in_config: + return legacy_params.get(key1, {}).get(key2, default_val) + return getattr(self, trait_name) + + self.type_checkpoint = get_conf_val("data", "type_checkpoint", "type_checkpoint", self.type_checkpoint) + self.energy_checkpoint = get_conf_val("data", "energy_checkpoint", "energy_checkpoint", self.energy_checkpoint) + self.direction_checkpoint = get_conf_val("data", "direction_checkpoint", "direction_checkpoint", self.direction_checkpoint) + self.load_onnx_model = get_conf_val("data", "load_onnx_model", "load_onnx_model", self.load_onnx_model) + + self.experiment_number = get_conf_val("run_details", "experiment_number", "experiment_number", self.experiment_number) + self.save_onnx = get_conf_val("run_details", "save_onnx", "save_onnx", self.save_onnx) + self.save_k_checkpoints = get_conf_val("hyp", "save_k", "save_k_checkpoints", self.save_k_checkpoints) + self.device_str = get_conf_val("arch", "device", "device", self.device) + self.batch_size = get_conf_val("hyp", "batches", "batch_size", self.batch_size) + self.pin_memory = get_conf_val("dataset", "pin_memory", "pin_memory", self.pin_memory) + self.num_workers = get_conf_val("dataset", "num_workers", "num_workers", self.num_workers) + self.persistent_workers = get_conf_val("dataset", "persistent_workers", "persistent_workers", self.persistent_workers) + self.devices = get_conf_val("arch", "devices", "devices", self.devices) + self.strategy = get_conf_val("arch", "strategy", "strategy", self.strategy) + + # Augmentations + self.use_augmentation = get_conf_val("augmentation", "use_augmentation", "use_augmentation", self.use_augmentation) + self.aug_prob = get_conf_val("augmentation", "aug_prob", "aug_prob", self.aug_prob) + self.rot_prob = get_conf_val("augmentation", "rot_prob", "rot_prob", self.rot_prob) + self.trans_prob = get_conf_val("augmentation", "trans_prob", "trans_prob", self.trans_prob) + self.flip_hor_prob = get_conf_val("augmentation", "flip_hor_prob", "flip_hor_prob", self.flip_hor_prob) + self.flip_ver_prob = get_conf_val("augmentation", "flip_ver_prob", "flip_ver_prob", self.flip_ver_prob) + self.mask_prob = get_conf_val("augmentation", "mask_prob", "mask_prob", self.mask_prob) + self.mask_dvr_prob = get_conf_val("augmentation", "mask_dvr_prob", "mask_dvr_prob", self.mask_dvr_prob) + self.noise_prob = get_conf_val("augmentation", "noise_prob", "noise_prob", self.noise_prob) + self.max_rot = get_conf_val("augmentation", "max_rot", "max_rot", self.max_rot) + self.max_trans = get_conf_val("augmentation", "max_trans", "max_trans", self.max_trans) + + # Normalizations + self.use_clean = get_conf_val("normalization", "use_clean", "use_clean", self.use_clean) + self.use_clean_dvr = get_conf_val("normalization", "use_clean_dvr", "use_clean_dvr", self.use_clean_dvr) + self.type_mu = get_conf_val("normalization", "type_mu", "type_mu", self.type_mu) + self.type_sigma = get_conf_val("normalization", "type_sigma", "type_sigma", self.type_sigma) + self.dir_mu = get_conf_val("normalization", "dir_mu", "dir_mu", self.dir_mu) + self.dir_sigma = get_conf_val("normalization", "dir_sigma", "dir_sigma", self.dir_sigma) + self.energy_mu = get_conf_val("normalization", "energy_mu", "energy_mu", self.energy_mu) + self.energy_sigma = get_conf_val("normalization", "energy_sigma", "energy_sigma", self.energy_sigma) + + # Cut-offs + self.leakage_intensity_cutoff = get_conf_val("cut-off", "leakage_intensity", "leakage_intensity_cutoff", self.leakage_intensity_cutoff) + self.intensity_cutoff = get_conf_val("cut-off", "intensity", "intensity_cutoff", self.intensity_cutoff) + + self.hyp_configs = legacy_params.get("hyp", {}) + else: + self.log.info("No legacy config file provided. Using standard Traitlets configuration.") + self.device_str = self.device + self.hyp_configs = { + "epochs": self.n_epochs, + "batches": self.batch_size, + "dynamic_batches": True, + "optimizer": self.optimizer.get("name", "Adam"), + "momentum": self.optimizer_momentum, + "weight_decay": self.optimizer_weight_decay, + "learning_rate": self.optimizer.get("base_learning_rate", 0.0001), + "lrf": self.lrf, + "start_epoch": 0, + "steps_epoch": 100, + "l2_lambda": self.l2_lambda, + "adam_epsilon": self.optimizer.get("adam_epsilon", 1.0e-8), + "gradient_clip_val": self.gradient_clip_val, + "save_k": self.save_k_checkpoints, + } + + self.save_k = self.save_k_checkpoints + + self.parameters = { + "data": { + "train_gamma_proton": None, + "validation_gamma_proton": None, + "train_gamma": None, + "validation_gamma": None, + "test_gamma": None, + "test_proton": None, + "test_electron": None, + "test_validation_gamma": None, + "test_validation_gamma_proton": None, + "type_checkpoint": self.type_checkpoint, + "energy_checkpoint": self.energy_checkpoint, + "direction_checkpoint": self.direction_checkpoint, + "load_onnx_model": self.load_onnx_model, + }, + "run_details": { + "mode": "train", + "task": self.reco_tasks[0] if self.reco_tasks else "all", + "test_type": "gamma", + "experiment_number": self.experiment_number, + }, + "cut-off": { + "leakage_intensity": self.leakage_intensity_cutoff, + "intensity": self.intensity_cutoff, + }, + "model": self.model_type, + "hyp": self.hyp_configs, + "augmentation": { + "use_augmentation": self.use_augmentation, + "aug_prob": self.aug_prob, + "rot_prob": self.rot_prob, + "trans_prob": self.trans_prob, + "flip_hor_prob": self.flip_hor_prob, + "flip_ver_prob": self.flip_ver_prob, + "mask_prob": self.mask_prob, + "mask_dvr_prob": self.mask_dvr_prob, + "noise_prob": self.noise_prob, + "max_rot": self.max_rot, + "max_trans": self.max_trans, + }, + "normalization": { + "apply_log_scaling": self.apply_log_scaling, + "use_clean": self.use_clean, + "use_clean_dvr": self.use_clean_dvr, + "type_mu": self.type_mu, + "type_sigma": self.type_sigma, + "dir_mu": self.dir_mu, + "dir_sigma": self.dir_sigma, + "energy_mu": self.energy_mu, + "energy_sigma": self.energy_sigma, + }, + "dataset": { + "num_workers": self.num_workers, + "pin_memory": self.pin_memory, + "persistent_workers": self.persistent_workers, + }, + "arch": { + "device": self.device_str, + "precision_type": self.precision_type, + "precision_energy": self.precision_energy, + "precision_direction": self.precision_direction, + "devices": self.devices, + "strategy": self.strategy, + } + } + + print(f"Using Devices: {self.devices}") + + # Set up the data loaders for training and validation + indices = list(range(self.dl1dh_reader._get_n_events())) + # Shuffle the indices before the training/validation split + np.random.seed(self.random_seed) + np.random.shuffle(indices) + n_validation_examples = int( + self.validation_split * self.dl1dh_reader._get_n_events() + ) + training_indices = indices[n_validation_examples:] + validation_indices = indices[:n_validation_examples] + + if not ("class_weight" in self.parameters): + self.parameters['class_weight'] = self.dl1dh_reader.class_weight + self.log.info(f"Class weights not provided. Using class weights from data reader: {self.parameters['class_weight']}") + elif len(self.parameters['class_weight']) != len(self.dl1dh_reader.class_names): + raise ValueError(f"Number of class weights provided ({len(self.parameters['class_weight'])}) does not match number of classes in data ({len(self.dl1dh_reader.class_names)}).") + else: + self.log.info(f"Using class weights from configuration file: {self.parameters['class_weight']}") + + print("BASE TRAIN FRAMEWORK", self.framework_type) + + self.train_dataset = DLDataLoader.create( + framework=self.framework_type, + DLDataReader=self.dl1dh_reader, + indices=training_indices, + tasks=self.reco_tasks, + batch_size=self.batch_size, + random_seed=self.random_seed, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + parameters=self.parameters, + use_augmentation=self.use_augmentation, + is_training=True, + ) + self.training_loader = DataLoader( + dataset=self.train_dataset, + batch_size=None, + batch_sampler=None, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + prefetch_factor=4 if self.num_workers > 0 else None, + persistent_workers=self.persistent_workers if self.num_workers > 0 else False + ) + + print(len(self.training_loader)) + + self.validation_dataset = DLDataLoader.create( + framework=self.framework_type, + DLDataReader=self.dl1dh_reader, + indices=validation_indices, + tasks=self.reco_tasks, + batch_size=self.batch_size, + random_seed=self.random_seed, + sort_by_intensity=self.sort_by_intensity, + stack_telescope_images=self.stack_telescope_images, + parameters=self.parameters, + use_augmentation=False, + is_training=False, + ) + self.validation_loader = DataLoader( + dataset=self.validation_dataset, + batch_size=None, + batch_sampler=None, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + prefetch_factor=4 if self.num_workers > 0 else None, + persistent_workers=self.persistent_workers if self.num_workers > 0 else False + ) + + print(len(self.validation_dataset)) + + + def start(self): + super().start() + for task in self.tasks: + + # Create the experiment folder + save_folder = create_experiment_folder( + f"run_{task.name}_training_", next_number=self.experiment_number + ) + + # Save the resolved configuration parameters to a config file in the experiment directory + config_filename = os.path.join(save_folder, "resolved_config.yml") + with open(config_filename, "w") as outfile: + import yaml + yaml.dump(self.parameters, outfile, default_flow_style=False) + + # ------------------------------------------------------------------------------ + # Select the model and precision + # ------------------------------------------------------------------------------ + import torch.nn as nn + import torch + torch.backends.cudnn.enabled = False + + if task == Task.type: + precision = self.parameters["arch"]["precision_type"] + elif task == Task.energy: + precision = self.parameters["arch"]["precision_energy"] + elif task == Task.cameradirection or task == Task.skydirection: + precision = self.parameters["arch"]["precision_direction"] + else: + raise ValueError( + f"task:{task.name} is not supported. Task must be type, direction or energy" + ) + + class ONNXModelWrapper(nn.Module): + def __init__(self, onnx_model_net, active_task, onnx_input_shape): + super().__init__() + self.onnx_model_net = onnx_model_net + self.active_task = active_task + self.onnx_input_shape = onnx_input_shape + + self.onnx_channel_last = False + self.expected_channels = 1 + if len(onnx_input_shape) == 4: + if onnx_input_shape[3] in [1, 2, 3] and onnx_input_shape[1] > onnx_input_shape[3]: + self.onnx_channel_last = True + self.expected_channels = onnx_input_shape[3] + else: + self.expected_channels = onnx_input_shape[1] + + def forward(self, x, y=None): + if self.expected_channels == 2 and y is not None: + x = torch.cat([x, y], dim=1) + + import inspect + sig = inspect.signature(self.onnx_model_net.forward) + num_onnx_inputs = len(sig.parameters) + + if num_onnx_inputs == 2 and y is not None: + if self.onnx_channel_last: + x = x.permute(0, 2, 3, 1) + y = y.permute(0, 2, 3, 1) + out = self.onnx_model_net(x, y) + else: + if self.onnx_channel_last: + x = x.permute(0, 2, 3, 1) + out = self.onnx_model_net(x) + + if isinstance(out, (tuple, list)): + if len(out) == 3: + return out + val = out[0] + else: + val = out + + if self.active_task == Task.type: + return val, None, None + elif self.active_task == Task.energy: + return None, val, None + else: + return None, None, val + + if self.load_onnx_model: + self.log.info(f"Loading ONNX model from {self.load_onnx_model} for training...") + self.log.warning( + "WARNING: Training an ONNX model from scratch (untrained) in PyTorch using onnx2pytorch " + "often leads to frozen gradients and the loss not improving. This is because onnx2pytorch " + "is designed primarily for inference, and many operations break the computational graph. " + "It is highly recommended to train the native PyTorch model instead by removing the " + "--load_onnx_model flag." + ) + import onnx + from onnx2pytorch import ConvertModel + try: + onnx_proto = onnx.load(self.load_onnx_model) + onnx_model = ConvertModel(onnx_proto) + onnx_input_shape = [dim.dim_value for dim in onnx_proto.graph.input[0].type.tensor_type.shape.dim] + model_net = ONNXModelWrapper(onnx_model, task, onnx_input_shape) + except Exception as e: + self.log.error(f"Failed to load ONNX model: {e}") + raise e + else: + self.log.info("Setting up the PyTorch model.") + num_inputs = 1 + if isinstance(self.parameters.get("model"), dict): + num_inputs = self.parameters["model"].get(f"model_{task.name}", {}).get("parameters", {}).get("num_inputs", 1) + + model_input_shape = list(self.train_dataset.input_shape) + if len(model_input_shape) == 3: + # Convert (H, W, C) to PyTorch's (C, H, W) + model_input_shape = [model_input_shape[2], model_input_shape[0], model_input_shape[1]] + elif len(model_input_shape) >= 4: + # e.g., (Telescopes, H, W, C) -> (Telescopes, C, H, W) + model_input_shape = list(model_input_shape[:-3]) + [model_input_shape[-1], model_input_shape[-3], model_input_shape[-2]] + + model_net = CTLearnModel.from_name( + self.model_type, + input_shape=tuple(model_input_shape), + tasks=[task.name], + parent=self, + ).model + + # if hasattr(model_net, 'T'): + # self.training_loader.set_T(model_net.T) + # ------------------------------------------------------------------------------ + # Load Checkpoints + # ------------------------------------------------------------------------------ + if task == Task.type: + check_point_path = self.parameters["data"]["type_checkpoint"] + + elif task == Task.energy: + check_point_path = self.parameters["data"]["energy_checkpoint"] + + elif task == Task.cameradirection or task == Task.skydirection: + check_point_path = self.parameters["data"]["direction_checkpoint"] + + else: + raise ValueError( + f"task:{task.name} is not supported. Task must be type, direction or energy" + ) + # Load the checkpoint if provided + if check_point_path: + model_net = ModelHelper.loadModel( + model_net, "", check_point_path, Mode.train, device_str=self.device_str + ) + + # Setup the TensorBoard logger + log_dir = save_folder + + tb_logger = TensorBoardLogger( + save_dir=log_dir, + name="exp_" + + str(self.experiment_number) + + "_" + + task.name + + "_train", + default_hp_metric=False, + ) + + extra_trainer_args = {} + test_limit = os.environ.get("CTLEARN_TEST_LIMIT") + if test_limit: + try: + limit = int(test_limit) + except ValueError: + limit = 5 + extra_trainer_args["limit_train_batches"] = limit + extra_trainer_args["limit_val_batches"] = limit + + # Setup the Trainer + trainer_pl = CTLearnTrainer( + max_epochs=self.parameters["hyp"]["epochs"], + accelerator=self.parameters["arch"]["device"], + devices=self.devices, + strategy= self.parameters["arch"]["strategy"], + default_root_dir=log_dir, + log_every_n_steps=1, + logger=tb_logger, + num_sanity_val_steps=0, + precision=precision, + gradient_clip_val=self.parameters["hyp"]["gradient_clip_val"], + callbacks=[GPUStatsLogger()], + sync_batchnorm=True, + enable_progress_bar=not self.disable_progress_bar, + **extra_trainer_args + ) + + # Setup Lighting + lightning_model = CTLearnPL( + model=model_net, + save_folder=trainer_pl.get_log_dir(), + task=task, + mode = Mode.train, + parameters=self.parameters, + k=self.save_k, + train_loader= self.training_loader, + val_loader= self.validation_loader, + train_dataset=self.train_dataset, + val_dataset=self.validation_dataset + ) + + if trainer_pl.is_global_zero: + # Save configuration file. + if not os.path.exists(trainer_pl.get_log_dir()): + os.makedirs(trainer_pl.get_log_dir()) + + with open(os.path.join(trainer_pl.get_log_dir(),"parameters.json"), "w") as f: + def path_serializer(obj): + import pathlib + if isinstance(obj, pathlib.Path): + return str(obj) + raise TypeError(f"Type {type(obj)} not serializable") + json.dump(self.parameters, f, indent=4, default=path_serializer) + + print(f"Run tensorboard server: tensorboard --load_fast=false --host=0.0.0.0 --logdir={trainer_pl.get_log_dir()}/") + + print(f"Accelerator: {trainer_pl.accelerator}") + print(f"Num. Devices: {trainer_pl.num_devices}") + + trainer_pl.fit( + model=lightning_model, + train_dataloaders=self.training_loader, + val_dataloaders=[self.validation_loader], + ) + + # Export to ONNX if requested + if self.save_onnx: + self.log.info("Converting PyTorch model into ONNX format...") + try: + # Load the best model weights if available + if trainer_pl.checkpoint_callback and trainer_pl.checkpoint_callback.best_model_path: + best_model_path = trainer_pl.checkpoint_callback.best_model_path + self.log.info(f"Loading best model from {best_model_path} for ONNX export...") + model_net = ModelHelper.loadModel( + model_net, "", best_model_path, Mode.train, device_str=self.device_str + ) + + # Create dummy input dynamically from a batch + batch = next(iter(self.training_loader)) + features, labels, t = batch + + import inspect + sig = inspect.signature(model_net.forward) + num_inputs_sig = len(sig.parameters) + + # Transfer dummy inputs to same device as model + if num_inputs_sig == 2: + dummy_image = torch.randn_like(features["image"][:1]).to(self.device_str) + dummy_peak = torch.randn_like(features["peak_time"][:1]).to(self.device_str) + dummy_input = (dummy_image, dummy_peak) + input_names = ["image", "peak_time"] + else: + dummy_input = torch.randn_like(features["image"][:1]).to(self.device_str) + input_names = ["image"] + + output_names = [task.name] + + # Put model in evaluation mode + model_net.eval() + + onnx_path = os.path.join(save_folder, f"ctlearn_model_{task.name}") + ModelHelper.exportOnnx( + model_net, + dummy_input, + onnx_path, + input_names=input_names, + output_names=output_names + ) + self.log.info(f"ONNX model saved successfully to {onnx_path}.onnx and {onnx_path}_simp.onnx") + except Exception as e: + self.log.error(f"Failed to export model to ONNX: {e}") + + def finish(self): + super().finish() + print("Pytorch finish") + + def show_version(self): + print("Pytorch 2.3") diff --git a/ctlearn/tools/train/pytorch/utils.py b/ctlearn/tools/train/pytorch/utils.py new file mode 100644 index 00000000..4d10fe47 --- /dev/null +++ b/ctlearn/tools/train/pytorch/utils.py @@ -0,0 +1,190 @@ +from ctlearn.core.ctlearn_enum import Task +from typing import List +import os +import yaml + +expected_structure = { + "data": { + "train_gamma_proton": None, + "validation_gamma_proton": None, + "train_gamma": None, + "validation_gamma": None, + "test_gamma": None, + "test_proton": None, + "test_electron": None, + "test_validation_gamma": None, + "test_validation_gamma_proton": None, + }, + "run_details": { + "mode": None, + "task": None, + "test_type": None, + "experiment_number": None, + }, + "cut-off": {"leakage_intensity": None, "intensity": None}, + "model": { + "model_type": { + "model_name": None, + "parameters": None, + }, + "model_energy": { + "model_name": None, + "parameters": None, + }, + "model_direction": { + "model_name": None, + "parameters": None, + }, + }, + "hyp": { + "epochs": None, + "batches": None, + "dynamic_batches": None, + "optimizer": None, + "momentum": None, + "weight_decay": None, + "learning_rate": None, + "lrf": None, + "start_epoch": None, + "steps_epoch": None, + "l2_lambda": None, + "adam_epsilon": None, + "gradient_clip_val": None, + "save_k": None, + }, + "augmentation": { + "use_augmentation": None, + "aug_prob": None, + "rot_prob": None, + "trans_prob": None, + "flip_hor_prob": None, + "flip_ver_prob": None, + "mask_prob": None, + "noise_prob": None, + "max_rot": None, + "max_trans": None, + }, + "normalization": { + "use_clean": None, + "type_mu": None, + "type_sigma": None, + "dir_mu": None, + "dir_sigma": None, + "energy_mu": None, + "energy_sigma": None, + }, + "dataset": {"num_workers": None, "pin_memory": None, "persistent_workers": None}, + "arch": { + "device": None, + "precision_type": None, + "precision_energy": None, + "precision_direction": None, + "devices": None, + "strategy": None, + }, +} + + +# ------------------------------------------------------------------------------------------------------------------- +# Sanity check function +def sanity_check(config, expected_structure): + """ + Recursively checks if all required keys in the expected_structure exist in the config data. + Raises a KeyError if a key is missing. + """ + for key, substructure in expected_structure.items(): + if key not in config: + raise KeyError(f"Missing key: {key}") + + # If the substructure is a dictionary, recursively check the subkeys + if isinstance(substructure, dict): + if not isinstance(config[key], dict): + raise KeyError( + f"Expected a dictionary for key: {key}, but got: {type(config[key])}" + ) + sanity_check(config[key], substructure) + + +# ------------------------------------------------------------------------------------------------------------------- +def read_configuration(config_file_str="./config/training_config.yml"): + parameters = None + + if os.path.exists(config_file_str): + with open(config_file_str, "r") as config_file: + parameters = yaml.safe_load(config_file) + + else: + print(f"Configuration file not found. ({config_file_str})") + + return parameters + + +# ------------------------------------------------------------------------------------------------------------------- +def create_experiment_folder(prefix="run_", next_number=None): + """ + Create the next folder within the specified directory with a given prefix. + If next_number is not specified, automatically find the next available number. + + :param run_directory: The directory where folders are managed. + :param prefix: Prefix used for folders. + :param next_number: Optional. Specify the number to be used for the new folder. + """ + run_directory = "./run" + + # Ensure the 'run' directory exists + if not os.path.exists(run_directory): + os.makedirs(run_directory) + print(f"Directory '{run_directory}' created.") + + if next_number is None: + # List all subdirectories in the 'run' directory + folders = [ + f + for f in os.listdir(run_directory) + if os.path.isdir(os.path.join(run_directory, f)) + ] + # Filter folders that match the prefix pattern and end with a digit + matching_folders = [ + f for f in folders if f.startswith(prefix) and f[len(prefix) :].isdigit() + ] + + # Find the highest number and calculate the next one + if matching_folders: + highest_number = max( + int(folder[len(prefix) :]) for folder in matching_folders + ) + next_number = highest_number + 1 + else: + next_number = 0 + + # Create the new folder with the next number + new_folder_name = f"{prefix}{next_number}" + new_folder_path = os.path.join(run_directory, new_folder_name) + new_folder_path += "/" + if not os.path.exists(new_folder_path): + os.makedirs(new_folder_path) + print(f"New folder created: {new_folder_path}") + return new_folder_path + + +# ------------------------------------------------------------------------------------------------------------------- +def str_list_to_enum_list(reco_tasks: List) -> List[Task]: + + tasks = [] + for task_str in reco_tasks: + try: + tasks.append(Task[task_str]) + except KeyError: + print(f"'{task_str}' is not a valid enum type.") + return tasks + + +# ------------------------------------------------------------------------------------------------------------------- +def get_absolute_config_path(config_file_str="./ctlearn/tools/train/pytorch/config/default_config_file.yml") -> str: + """ + Returns the absolute path of the configuration file. + + :param config_file_str: Relative or absolute path to the configuration file. + :return: Absolute path of the configuration file. + """ + return os.path.abspath(config_file_str) \ No newline at end of file diff --git a/ctlearn/tools/train_model.py b/ctlearn/tools/train_model.py index e4f729d2..19be869d 100644 --- a/ctlearn/tools/train_model.py +++ b/ctlearn/tools/train_model.py @@ -1,230 +1,66 @@ -""" -Tool to train a ``CTLearnModel`` on R1/DL1a data using the ``DLDataReader`` and ``DLDataLoader``. -""" +import sys +import os +import warnings + +is_debug = '--debug' in sys.argv or any(arg.startswith('--log-level=DEBUG') for arg in sys.argv) +if not is_debug: + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' + os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' + os.environ['NCCL_DEBUG'] = 'WARN' + warnings.filterwarnings("ignore", category=UserWarning) + warnings.filterwarnings("ignore", ".*NoneDefaultNotAllowedWarning.*") + warnings.filterwarnings("ignore", ".*MergeConflictWarning.*") + warnings.filterwarnings("ignore", ".*'ctlearn.tools.train_model' found in sys.modules.*") import atexit -import keras import pandas as pd import numpy as np -import tensorflow as tf from ctapipe.core import Tool -from ctapipe.core.tool import ToolConfigurationError -from ctapipe.core.traits import ( - Bool, - CaselessStrEnum, - Path, - Float, - Int, - List, - Dict, - classes_with_traits, - ComponentName, - Unicode, -) -from 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 - - -class TrainCTLearnModel(Tool): +from ctapipe.core.traits import CaselessStrEnum, Dict +from ctlearn.core.ctlearn_enum import FrameworkType +class DLFrameWork(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 + Tool to select and run a specific deep learning training framework (Keras or PyTorch) + for CTLearn model training. It dynamically loads the appropriate subclass based on + the user-defined --framework argument. """ - name = "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 \\ - """ - - input_dir_signal = Path( - help="Input directory for signal events", - allow_none=False, - exists=True, - directory_ok=True, - file_ok=False, - ).tag(config=True) + name = "dlframework" - file_pattern_signal = List( - trait=Unicode(), - default_value=["*.h5"], - help="List of specific file pattern for matching files in ``input_dir_signal``", + framework_type = CaselessStrEnum( + ["pytorch", "keras"], + default_value="keras", + help="Framework to use: pytorch or keras", ).tag(config=True) - input_dir_background = Path( + early_stopping = Dict( default_value=None, - help="Input directory for background events", allow_none=True, - exists=True, - directory_ok=True, - file_ok=False, - ).tag(config=True) - - file_pattern_background = List( - trait=Unicode(), - default_value=["*.h5"], - help="List of specific file pattern for matching files in ``input_dir_background``", - ).tag(config=True) - - dl1dh_reader_type = ComponentName(DLDataReader, default_value="DLImageReader").tag( - config=True - ) - - stack_telescope_images = Bool( - default_value=False, - allow_none=False, help=( - "Set whether to stack the telescope images in the data loader. " - "Requires DLDataReader mode to be ``stereo``." + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " ), ).tag(config=True) - sort_by_intensity = Bool( - default_value=False, + early_stopping = Dict( + default_value=None, allow_none=True, help=( - "Set whether to sort the telescope images by intensity in the data loader. " - "Requires DLDataReader mode to be ``stereo``." + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " ), ).tag(config=True) - model_type = ComponentName(CTLearnModel, default_value="ResNet").tag(config=True) - - output_dir = Path( - exits=False, + early_stopping = Dict( default_value=None, - allow_none=False, - directory_ok=True, - file_ok=False, - help="Output directory for the trained reconstructor.", - ).tag(config=True) - - reco_tasks = List( - trait=CaselessStrEnum(["type", "energy", "cameradirection", "skydirection"]), - allow_none=False, - help=( - "List of reconstruction tasks to perform. " - "'type': classification of the primary particle type; " - "'energy': regression of the primary particle energy; " - "'cameradirection': reconstruction of the primary particle arrival direction in camera coordinates; " - "'skydirection': reconstruction of the primary particle arrival direction in sky coordinates." - ), - ).tag(config=True) - - n_epochs = Int( - default_value=10, - allow_none=False, - help="Number of epochs to train the neural network.", - ).tag(config=True) - - batch_size = Int( - default_value=64, - allow_none=False, - help="Size of the batch to train the neural network.", - ).tag(config=True) - - validation_split = Float( - default_value=0.1, - help="Fraction of the data to use for validation", - min=0.01, - max=0.99, - ).tag(config=True) - - 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", - "base_learning_rate": 0.0001, - "adam_epsilon": 1.0e-8, - }, - help=( - "Optimizer to use for training. " - "E.g. {'name': 'Adam', 'base_learning_rate': 0.0001, 'adam_epsilon': 1.0e-8}. " - ), - ).tag(config=True) - - 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=( - "Random seed for shuffling the data " - "before the training/validation split " - "and after the end of an epoch." + "Early stopping parameters for the Keras callback. " + "E.g. {'monitor': 'val_loss', 'patience': 4, 'verbose': 1, 'restore_best_weights': True}. " ), ).tag(config=True) - save_onnx = Bool( - default_value=False, - allow_none=False, - help="Set whether to save model in an ONNX file.", - ).tag(config=True) - early_stopping = Dict( default_value=None, allow_none=True, @@ -235,321 +71,204 @@ class TrainCTLearnModel(Tool): ).tag(config=True) aliases = { - "signal": "TrainCTLearnModel.input_dir_signal", - "background": "TrainCTLearnModel.input_dir_background", - "pattern-signal": "TrainCTLearnModel.file_pattern_signal", - "pattern-background": "TrainCTLearnModel.file_pattern_background", - "reco": "TrainCTLearnModel.reco_tasks", - ("o", "output"): "TrainCTLearnModel.output_dir", + "framework": "DLFrameWork.framework_type", } - - classes = classes_with_traits(CTLearnModel) + classes_with_traits(DLDataReader) + + try: + from ctlearn.tools.train.pytorch.train_pytorch_model import TrainPyTorchModel + aliases.update(TrainPyTorchModel.aliases) + except ImportError: + pass + + try: + from ctlearn.tools.train.keras.train_keras_model import TrainKerasModel + aliases.update(TrainKerasModel.aliases) + except ImportError: + pass + + def __init__(self, **kwargs): + """ + Initialize the DLFrameWork tool and prepare for framework injection. + """ + super().__init__(**kwargs) + self.framework_instance = None def setup(self): - self.log.info("ctlearn version %s", ctlearn_version) - # Check if the output directory exists - if self.output_dir.exists(): - raise ToolConfigurationError( - f"Output directory {self.output_dir} already exists." - ) - # 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: - self.input_url_signal.extend(self.input_dir_signal.glob(signal_pattern)) - # Get bkg input files - self.input_url_background = [] - if self.input_dir_background is not None: - for background_pattern in self.file_pattern_background: - self.input_url_background.extend( - self.input_dir_background.glob(background_pattern) - ) + """ + Setup method called after basic trait parsing. + This dynamically loads and prepares the correct framework subclass + (TrainKerasModel or TrainPyTorchModel). + """ + framework_enum = self.string_to_type(self.framework_type) + self.framework_instance = self.get_framework(framework_enum) - # Set up the data reader - self.log.info("Loading data:") - self.log.info("For a large dataset, this may take a while...") - if self.dl1dh_reader_type == "DLFeatureVectorReader": - raise NotImplementedError( - "'DLFeatureVectorReader' is not supported in CTLearn yet. " - "Missing stereo CTLearnModel implementation." - ) - self.dl1dh_reader = DLDataReader.from_name( - self.dl1dh_reader_type, - input_url_signal=sorted(self.input_url_signal), - input_url_background=sorted(self.input_url_background), - parent=self, - ) - self.log.info("Number of events loaded: %s", self.dl1dh_reader._get_n_events()) - if "type" in self.reco_tasks: - self.log.info( - "Number of signal events: %d", self.dl1dh_reader.n_signal_events - ) - self.log.info( - "Number of background events: %d", self.dl1dh_reader.n_bkg_events - ) - # Check if the number of events is enough to form a batch - if self.dl1dh_reader._get_n_events() < self.batch_size: - raise ValueError( - f"{self.dl1dh_reader._get_n_events()} events are not enough " - f"to form a batch of size {self.batch_size}. Reduce the batch size." - ) - # Check if there are at least two classes in the reader for the particle classification - if self.dl1dh_reader.class_weight is None and "type" in self.reco_tasks: - raise ValueError( - "Classification task selected but less than two classes are present in the data." - ) - # Check if stereo mode is selected for stacking telescope images - if self.stack_telescope_images and self.dl1dh_reader.mode == "mono": - raise ToolConfigurationError( - f"Cannot stack telescope images in mono mode. Use stereo mode for stacking." - ) - # Ckeck if only one telescope type is selected for stacking telescope images - if ( - self.stack_telescope_images - and len(list(self.dl1dh_reader.selected_telescopes)) > 1 - ): - raise ToolConfigurationError( - f"Cannot stack telescope images from multiple telescope types. Use only one telescope type." - ) - # Check if sorting by intensity is disabled for stacking telescope images - if self.stack_telescope_images and self.sort_by_intensity: - raise ToolConfigurationError( - f"Cannot stack telescope images when sorting by intensity. Disable sorting by intensity." - ) - - # Set up the data loaders for training and validation - indices = list(range(self.dl1dh_reader._get_n_events())) - # Shuffle the indices before the training/validation split - np.random.seed(self.random_seed) - np.random.shuffle(indices) - n_validation_examples = int( - self.validation_split * self.dl1dh_reader._get_n_events() - ) - training_indices = indices[n_validation_examples:] - validation_indices = indices[:n_validation_examples] - 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, - ) - - # 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 - 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) + # Inject aliases and shared config before full CLI parsing + self.framework_instance.update_config(self.config) + self.aliases.update(self.framework_instance.aliases) + DLFrameWork.aliases.update(self.framework_instance.aliases) def start(self): + """ + Start method called after setup. Executes the selected framework instance. + """ + print("start") + self.framework_instance.run() - # 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!") - - 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!") + @classmethod + def string_to_type(cls, str_type: str) -> FrameworkType: + """ + Convert a string to a FrameworkType enum (case-insensitive). - 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) + Parameters: + str_type (str): The name of the framework (e.g., 'keras', 'pytorch'). - self.log.info("Tool is shutting down") + Returns: + FrameworkType: Corresponding enum value. - def _get_losses_and_mertics(self, tasks): + Raises: + ValueError: If the provided string is not a valid framework type. """ - 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. + try: + return FrameworkType[str_type.upper()] + except KeyError: + raise ValueError(f"'{str_type}' is not a valid framework type.") + + @classmethod + def get_framework(cls, framework_type: FrameworkType): + """ + Dynamically import and return the corresponding training class + based on the framework type. + + Parameters: + framework_type (FrameworkType): Enum indicating which framework to use. + + Returns: + Tool: An instance of the selected training framework (subclass of Tool). + + Raises: + ImportError: If the training module could not be imported. + ValueError: If the framework type is unknown. """ - 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 + if framework_type == FrameworkType.KERAS: + try: + from ctlearn.tools.train.keras.train_keras_model import TrainKerasModel + fw = TrainKerasModel() + except ImportError as e: + raise ImportError(f"Not possible to import TrainKerasModel: {e}") from e + + elif framework_type == FrameworkType.PYTORCH: + try: + from ctlearn.tools.train.pytorch.train_pytorch_model import ( + TrainPyTorchModel, + ) + + fw = TrainPyTorchModel() + + except ImportError as e: + raise ImportError( + f"Not possible to import TrainPyTorchModel: {e}" + ) from e + + else: + raise ValueError(f"Unknown Framework: {framework_type.name}") + + return fw + + from ctlearn.tools.train.base_train_model import TrainCTLearnModel + classes = [TrainCTLearnModel] + try: + from ctlearn.tools.train.keras.train_keras_model import TrainKerasModel + classes.append(TrainKerasModel) + except ImportError: + pass + + try: + from ctlearn.tools.train.pytorch.train_pytorch_model import TrainPyTorchModel + classes.append(TrainPyTorchModel) + except ImportError: + pass + + from ctapipe.core.traits import classes_with_traits + from dl1_data_handler.reader import DLDataReader + classes = classes + classes_with_traits(DLDataReader) def main(): # Run the tool - tool = TrainCTLearnModel() + tool = DLFrameWork() + + # Manually parse --framework to determine which subclass to load, as traitlets alias + # update can sometimes fail to parse it correctly before setup + framework = "keras" + for i, arg in enumerate(sys.argv[1:]): + if arg.startswith("--framework="): + framework = arg.split("=")[1].strip() + elif arg == "--framework" and i + 2 < len(sys.argv): + framework = sys.argv[i + 2].strip() + + tool.framework_type = framework + + minimal_args = [ + arg for arg in sys.argv[1:] if "--framework" in arg or arg in ["-h", "--help"] + ] + tool.initialize(argv=minimal_args) + + # Setup and inject the correct framework instance + tool.setup() + + # Parse all CLI args with the selected framework subclass + tool.framework_instance.initialize(argv=sys.argv[1:]) + tool.run() -if __name__ == "main": +if __name__ == "__main__": main() + + +# Example: +# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir2 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite +# python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal ./mc_tjark/ --pattern-signal gamma_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_*.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --pattern-background proton_theta_23.161_az_99.261_runs833-1250*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 --pattern-background proton_*.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 + +# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs1-60.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs181-240.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs188-246.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs247-305.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs1-59.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs177-235.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs181-240.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs181-240.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs181-240.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs1-60.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs181-240.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs1-60.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs181-240.dl1.h5 + + + + +# --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 + + +# --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 + + + +# Type +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --background /storage/ctlearn_data/h5_files/mc/protons/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --pattern-background=proton_theta_16.087_az_108.090_runs1-416.dl1.h5 --pattern-background=proton_theta_16.087_az_251.910_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_260.739_runs1-417.dl1.h5 --pattern-background=proton_theta_23.161_az_99.261_runs1-417.dl1.h5 --pattern-background=proton_theta_30.390_az_266.360_runs1-416.dl1.h5 --pattern-background=proton_theta_30.390_az_93.640_runs1-420.dl1.h5 --pattern-background=proton_theta_37.661_az_270.641_runs1-421.dl1.h5 --pattern-background=proton_theta_37.661_az_89.359_runs1-406.dl1.h5 --pattern-background=proton_theta_6.000_az_180.000_runs1-416.dl1.h5 --pattern-background=proton_theta_9.579_az_126.888_runs1-417.dl1.h5 --pattern-background=proton_theta_9.579_az_233.112_runs1-417.dl1.h5 --reco type --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_type_training.out 2>&1 & + +# Direction +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal gamma_theta_23.161_az_260.739_runs7-65*.dl1.h5 --reco cameradirection --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_direction_training.out 2>&1 & + + + +# Energy +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_251.910_runs121-180.dl1.h5 --pattern-signal=gamma_theta_23.161_az_260.739_runs129-187.dl1.h5 --pattern-signal=gamma_theta_23.161_az_99.261_runs118-176.dl1.h5 --pattern-signal=gamma_theta_30.390_az_266.360_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs121-180.dl1.h5 --pattern-signal=gamma_theta_30.390_az_93.640_runs1-60.dl1.h5 --pattern-signal=gamma_theta_37.661_az_270.641_runs121-180.dl1.h5 --pattern-signal=gamma_theta_37.661_az_89.359_runs121-180.dl1.h5 --pattern-signal=gamma_theta_6.000_az_180.000_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_126.888_runs121-180.dl1.h5 --pattern-signal=gamma_theta_9.579_az_233.112_runs121-180.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_1.yml> nohup_energy_training.out 2>&1 & + + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs243-302.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs303-362.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs363-422.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs423-482.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs483-541.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs542-600.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs63-122.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_6.000_az_*.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & + +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml> nohup_energy_training.out 2>&1 & +# nohup python -m ctlearn.tools.train_model --framework=pytorch --output ./output_dir3 --signal /storage/ctlearn_data/h5_files/mc/gamma-diffuse/ --pattern-signal=gamma_theta_16.087_az_108.090_runs123-182.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs1-62.dl1.h5 --pattern-signal=gamma_theta_16.087_az_108.090_runs183-242.dl1.h5 --reco energy --overwrite --config_file ./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training_v5_2.yml> nohup_energy_training.out 2>&1 & \ No newline at end of file diff --git a/ctlearn/utils.py b/ctlearn/utils.py index d2929527..ee16420e 100644 --- a/ctlearn/utils.py +++ b/ctlearn/utils.py @@ -2,31 +2,51 @@ from ctapipe.core import Provenance from ctapipe.core.traits import TraitError -from ctapipe.instrument import SubarrayDescription from ctapipe.instrument.optics import FocalLengthKind +from ctapipe.instrument import SubarrayDescription +import os +import time +from tqdm import tqdm -__all__ = ["get_lst1_subarray_description", "validate_trait_dict"] +__all__ = ["validate_trait_dict","get_lst1_subarray_description","monitor_progress"] -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``. +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 - 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) + 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. @@ -47,3 +67,22 @@ def validate_trait_dict(dict, required_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) \ No newline at end of file diff --git a/docs/source/usage.rst b/docs/source/usage.rst index c49050e1..f7f20e36 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: diff --git a/monitoring_gpus.py b/monitoring_gpus.py new file mode 100644 index 00000000..30c146e3 --- /dev/null +++ b/monitoring_gpus.py @@ -0,0 +1,108 @@ +import re +import requests +import paramiko +from dotenv import load_dotenv +import os +load_dotenv() + +# Configuration +FILE_URL = "https://www.cc.iaa.es/system/nodes/status" # Replace with the actual URL +LOCAL_FILE = "status.html" +SSH_USER = os.getenv("SSH_USER") +SSH_PASSWORD = os.getenv("SSH_PASSWORD") +COMMAND_TO_RUN = "conda activate ctlearn_pytorch && cd neutron_remote_connection && ls" # Replace with the actual command + +COMMANDS_FILE = "commands.txt" +STATE_FILE = "last_command_index.txt" + +def get_next_command(): + """Reads the next command to execute from the file, keeping track of progress.""" + if not os.path.exists(COMMANDS_FILE): + print("Commands file not found.") + return None + + with open(COMMANDS_FILE, "r") as f: + commands = f.readlines() + + if not commands: + print("No commands to execute.") + return None + + last_index = 0 + if os.path.exists(STATE_FILE): + with open(STATE_FILE, "r") as f: + last_index = int(f.read().strip()) + + if last_index >= len(commands): + print("All commands have been executed.") + return None + + next_command = commands[last_index].strip() + + with open(STATE_FILE, "w") as f: + f.write(str(last_index + 1)) + + return next_command + +def download_file(): + """ Downloads the HTML file using requests without SSL verification """ + try: + response = requests.get(FILE_URL, verify=False) # Disable SSL verification + with open(LOCAL_FILE, "wb") as file: + file.write(response.content) + print("File downloaded successfully (SSL verification disabled).") + except Exception as e: + print(f"Error downloading the file: {e}") + +def parse_gpu_usage(): + """ Parses the HTML file to find GPUs with usage below 20% """ + low_usage_nodes = [] + with open(LOCAL_FILE, "r", encoding="utf-8") as file: + content = file.read() + + # Regular expression to extract node names and GPU usage + pattern = re.compile( + r']*>.*?\s*' # Node (e.g., n1.iaa.es) + r']*>.*?\s*' # %CPU + r']*>.*?\s*' # Cores + r']*>.*?\s*' # %Mem + r']*>.*?\s*' # GBMem + r']*>([\d.]+)', # %GPU + re.DOTALL + ) + + for match in pattern.findall(content): + node, gpu_usage = match[0], float(match[1]) + if gpu_usage < 20: + low_usage_nodes.append(node) + + return low_usage_nodes + +def execute_remote_command(node): + """ Connects to the node via SSH and executes a command """ + command = get_next_command() + try: + print(f"Connecting to {node} to execute {command}...") + client = paramiko.SSHClient() + client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) + client.connect(node, username=SSH_USER, password=SSH_PASSWORD) + stdin, stdout, stderr = client.exec_command(command) + print(f"Command output on {node}:\n", stdout.read().decode()) + client.close() + except Exception as e: + print(f"SSH connection error on {node}: {e}") + +def main(): + download_file() + idle_gpus = parse_gpu_usage() + + if idle_gpus: + print(f"Detected GPUs with low usage on nodes: {idle_gpus}") + for node in idle_gpus: + execute_remote_command(node) + print(node) + else: + print("No GPUs with usage below 20% were found.") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index f2a5c057..c6e155ba 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.2.0", + "torchvision", + "torchmetrics", + "opencv-python", ] dynamic = ["version"] @@ -49,12 +54,33 @@ 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" @@ -62,8 +88,8 @@ documentation = "https://ctlearn.readthedocs.io/en/latest/" [project.scripts] ctlearn-train-model = "ctlearn.tools.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-mono-model = "ctlearn.tools.predict.predict_mono:main" +ctlearn-predict-stereo-model = "ctlearn.tools.predict_stereo:main" ctlearn-predict-LST1= "ctlearn.tools.predict_LST1:main" [tool.setuptools_scm] diff --git a/test_dataloader.py b/test_dataloader.py new file mode 100644 index 00000000..1112e947 --- /dev/null +++ b/test_dataloader.py @@ -0,0 +1,104 @@ +from ctlearn.core.data_loader.loader import DLDataLoader +from dl1_data_handler.reader import DLImageReader +from dl1_data_handler.reader import DLDataReader +import numpy as np +import matplotlib.pyplot as plt +from ctlearn.tools.train.pytorch.utils import read_configuration +def on_key(event): + # Check if the "Esc" key was pressed + if event.key == 'escape': + plt.close(event.canvas.figure) + exit() + + +# dl1_gamma_file = ["/storage/ctlearn_data/h5_files/mc/gamma_theta_16.087_az_108.090_runs123-182.r1.dl1.h5"] +dl1_gamma_file = ["/storage/ctlearn_data/h5_files/mc/gamma-diffuse/gamma_theta_9.579_az_233.112_runs481-540.dl1.h5"] +config_file = "./ctlearn/tools/train/pytorch/config/training_config_iaa_neutron_training.yml" +dl1dh_reader = DLDataReader.from_name( + "DLImageReader", + input_url_signal=dl1_gamma_file, + channels = ["cleaned_image","cleaned_peak_time"], + # input_url_background=sorted(self.input_url_background), + # parent=self, +) + +# dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config) +# dataloader = DLDataLoader.create("pytorch") +parameters = read_configuration(config_file) + +random_seed = 0 +indices = list(range(dl1dh_reader._get_n_events())) +training_loader = DLDataLoader.create( + framework="pytorch", + DLDataReader=dl1dh_reader, + indices=indices, + tasks=["type","energy","skydirection","cameradirection","hillas"], + batch_size=32, + random_seed=0, + sort_by_intensity=False, + stack_telescope_images=False, + parameters= parameters, + use_augmentation=True, + is_training=True, + +) + +for batch_idx, (features, labels,t ) in enumerate(training_loader): + + plt.rcParams['keymap.quit'].append(' ') + fig, axes = plt.subplots(1, 2, figsize=(15, 5)) + ax = axes.ravel() + # fig.canvas.mpl_connect('key_press_event', lambda evt: print(repr(evt.key))) + fig.canvas.mpl_connect('key_press_event', on_key) + + if len(features)>0: + for id in range(len(features["image"])): + # gammaness = 0 + # gammaness = labels["gammaness"][id] + + # if gammaness>0.9: + + image = features["image"][id] + # clean_image = features["clean_image"] + peak_time = features["peak_time"][id] + # clean_peak_time = features["clean_peak_time"] + labels_class = labels["type"][id] + + hillas_intensity = features["hillas"]["hillas_intensity"][id].cpu().detach().numpy() + leakage_pixels_width_1= features["hillas"]["leakage_pixels_width_1"][id].cpu().detach().numpy() + leakage_pixels_width_2= features["hillas"]["leakage_pixels_width_2"][id].cpu().detach().numpy() + + leakage_intensity_width_1 = features["hillas"]["leakage_intensity_width_1"][id].cpu().detach().numpy() + leakage_intensity_width_2 = features["hillas"]["leakage_intensity_width_2"][id].cpu().detach().numpy() + + morphology_n_islands= features["hillas"]["morphology_n_islands"][id].cpu().detach().numpy() + + if leakage_intensity_width_2<0.2: + for key in features["hillas"].keys(): + print(f"{key}: {features['hillas'][key][id].cpu().detach().numpy()}") + + image = np.transpose(image, (1, 2, 0)) + # clean_image = np.transpose(clean_image, (1, 2, 0)) + + peak_time = np.transpose(peak_time, (1, 2, 0)) + # clean_peak_time = np.transpose(clean_peak_time, (1, 2, 0)) + + ax[0].set_title( + f"\n qCharge:\n Hillas Intensity:{hillas_intensity} \n Leakage_p_w_2: {leakage_pixels_width_2} \n Leakage_i_w_2: {leakage_intensity_width_2} \n morphology_n_islands: {morphology_n_islands} \n labels_class: {str(labels_class)}" , fontsize=8 + ) + # ax[1].set_title( + # f"Peak time: \n Gammaness: {gammaness}" , fontsize=8 + # ) + ax[0].imshow(image, cmap="viridis") + ax[1].imshow(peak_time, cmap="viridis") + + # print(f"Gammaness: {gammaness}") + # if leakage_intensity_width_2<0.2: + # print(f"found {leakage_intensity_width_2}") + plt.show() + # plt.show(block=False) + # plt.pause(1) + # print(batch_idx) + plt.close() + + ii = 0 \ No newline at end of file