Merge pull request #94 from asogaard/pytorch-lightning

Implementing pytorch-lightning

Author: asogaard

.gitignore

@@ -134,3 +134,4 @@ docs/**/*.rst
 
 # Badges
 misc/badges
+lightning_logs/

envs/gnn_py38.yml

@@ -6,6 +6,6 @@ channels:
   - conda-forge
 dependencies:
   - python=3.8
-  - pytorch=1.9.1
+  - pytorch=1.9.0
   - cudatoolkit=11.1
   - pyg

examples/test_model_training_pytorch_lightning.py

@@ -0,0 +1,128 @@
+import logging
+import numpy as np
+import pandas as pd
+from timer import timer
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import EarlyStopping
+import torch
+from torch.optim.adam import Adam
+
+from gnn_reco.components.loss_functions import  LogCoshLoss, VonMisesFisher2DLoss
+from gnn_reco.components.utils import fit_scaler
+from gnn_reco.data.constants import FEATURES, TRUTH
+from gnn_reco.data.utils import get_equal_proportion_neutrino_indices
+from gnn_reco.models import Model
+from gnn_reco.models.detector.icecube import IceCube86
+from gnn_reco.models.gnn import DynEdge, ConvNet
+from gnn_reco.models.graph_builders import KNNGraphBuilder
+from gnn_reco.models.task.reconstruction import EnergyReconstruction
+from gnn_reco.models.training.callbacks import ProgressBar, PiecewiseLinearLR
+from gnn_reco.models.training.utils import get_predictions, make_train_validation_dataloader, save_results
+
+# Configurations
+timer.set_level(logging.INFO)
+logging.basicConfig(level=logging.INFO)
+torch.multiprocessing.set_sharing_strategy('file_system')
+
+# Constants
+features = FEATURES.ICECUBE86
+truth = TRUTH.ICECUBE86
+
+# Main function definition
+def main():
+
+    print(f"features: {features}")
+    print(f"truth: {truth}")
+
+    # Configuraiton
+    db = '/groups/icecube/leonbozi/datafromrasmus/GNNReco/data/databases/dev_level7_noise_muon_nu_classification_pass2_fixedRetro_v3/data/dev_level7_noise_muon_nu_classification_pass2_fixedRetro_v3.db'
+    pulsemap = 'SRTTWOfflinePulsesDC'
+    batch_size = 256
+    num_workers = 10
+    gpus = [0]
+    target = 'energy'
+    n_epochs = 5
+    patience = 5
+    archive = '/groups/icecube/asogaard/gnn/results'
+
+    # Common variables
+    train_selection, _ = get_equal_proportion_neutrino_indices(db)
+    train_selection = train_selection[0:50000]
+
+    training_dataloader, validation_dataloader = make_train_validation_dataloader(
+        db,
+        train_selection,
+        pulsemap,
+        features,
+        truth,
+        batch_size=batch_size,
+        num_workers=num_workers,
+    )
+
+    # Building model
+    detector = IceCube86(
+        graph_builder=KNNGraphBuilder(nb_nearest_neighbours=8),
+    )
+    gnn = DynEdge(
+        nb_inputs=detector.nb_outputs,
+    )
+    task = EnergyReconstruction(
+        hidden_size=gnn.nb_outputs,
+        target_label=target,
+        loss_function=LogCoshLoss(
+            transform_prediction_and_target=torch.log10,
+        ),
+    )
+    model = Model(
+        detector=detector,
+        gnn=gnn,
+        tasks=[task],
+        optimizer_class=Adam,
+        optimizer_kwargs={'lr': 1e-03, 'eps': 1e-03},
+        scheduler_class=PiecewiseLinearLR,
+        scheduler_kwargs={
+            'milestones': [0, len(training_dataloader) / 2, len(training_dataloader) * n_epochs],
+            'factors': [1e-2, 1, 1e-02],
+        },
+        scheduler_config={
+            'interval': 'step',
+        },
+     )
+
+    # Training model
+    callbacks = [
+        EarlyStopping(
+            monitor='val_loss',
+            patience=patience,
+        ),
+        ProgressBar(),
+    ]
+
+    trainer = Trainer(
+        gpus=gpus,
+        max_epochs=n_epochs,
+        callbacks=callbacks,
+        log_every_n_steps=1,
+    )
+
+    try:
+        trainer.fit(model, training_dataloader, validation_dataloader)
+    except KeyboardInterrupt:
+        print("[ctrl+c] Exiting gracefully.")
+        pass
+
+    # Saving predictions to file
+    results = get_predictions(
+        trainer,
+        model,
+        validation_dataloader,
+        [target + '_pred'],
+        [target, 'event_no'],
+    )
+
+    save_results(db, 'dynedge_energy_pytorch_lightning', results,archive, model)
+
+# Main function call
+if __name__ == "__main__":
+    main()

examples/test_model_training_sqlite.py

@@ -3,18 +3,17 @@
 
 import torch
 
-from gnn_reco.components.loss_functions import  VonMisesFisher2DLoss
+from gnn_reco.components.loss_functions import  LogCoshLoss, VonMisesFisher2DLoss
 from gnn_reco.components.utils import fit_scaler
 from gnn_reco.data.constants import FEATURES, TRUTH
 from gnn_reco.data.utils import get_equal_proportion_neutrino_indices
-from gnn_reco.legacy.reimplemented import LegacyVonMisesFisherLoss, LegacyAngularReconstruction
-from gnn_reco.models import Model
+from gnn_reco.legacy.callbacks import PiecewiseLinearScheduler
+from gnn_reco.legacy.trainers import Trainer, Predictor
+from gnn_reco.legacy.model import Model
 from gnn_reco.models.detector.icecube import IceCube86
 from gnn_reco.models.gnn import DynEdge, ConvNet
 from gnn_reco.models.graph_builders import KNNGraphBuilder
-from gnn_reco.models.task.reconstruction import AngularReconstructionWithKappa
-from gnn_reco.models.training.callbacks import PiecewiseLinearScheduler
-from gnn_reco.models.training.trainers import Trainer, Predictor
+from gnn_reco.models.task.reconstruction import EnergyReconstruction
 from gnn_reco.models.training.utils import make_train_validation_dataloader, save_results
 
 # Configurations
@@ -35,11 +34,11 @@ def main():
     # Configuraiton
     db = '/groups/icecube/leonbozi/datafromrasmus/GNNReco/data/databases/dev_level7_noise_muon_nu_classification_pass2_fixedRetro_v3/data/dev_level7_noise_muon_nu_classification_pass2_fixedRetro_v3.db'
     pulsemap = 'SRTTWOfflinePulsesDC'
-    batch_size = 1024
+    batch_size = 256
     num_workers = 10
-    device = 'cuda:1'
-    target = 'zenith'
-    n_epochs = 30
+    device = 'cuda:0'
+    target = 'energy'
+    n_epochs = 5
     patience = 5
     archive = '/groups/icecube/asogaard/gnn/results'
 
@@ -48,14 +47,14 @@ def main():
 
     # Common variables
     train_selection, _ = get_equal_proportion_neutrino_indices(db)
-    train_selection = train_selection[0:500000]
+    train_selection = train_selection[0:50000]
 
     training_dataloader, validation_dataloader = make_train_validation_dataloader(
-        db, 
-        train_selection, 
-        pulsemap, 
-        features, 
-        truth, 
+        db,
+        train_selection,
+        pulsemap,
+        features,
+        truth,
         batch_size=batch_size,
         num_workers=num_workers,
     )
@@ -68,19 +67,13 @@ def main():
     gnn = DynEdge(
         nb_inputs=detector.nb_outputs,
     )
-    task = LegacyAngularReconstruction(
-        hidden_size=gnn.nb_outputs, 
-        target_label=target, 
-        loss_function=LegacyVonMisesFisherLoss(
-            target_scaler=scalers['truth'][target]
+    task = EnergyReconstruction(
+        hidden_size=gnn.nb_outputs,
+        target_label=target,
+        loss_function=LogCoshLoss(
+            transform_prediction_and_target=torch.log10,
         ),
-        target_scaler=scalers['truth'][target],
     )
-    #task = AngularReconstructionWithKappa(
-    #    hidden_size=gnn.nb_outputs, 
-    #    target_label=target, 
-    #    loss_function=VonMisesFisher2DLoss(),
-    #)
     model = Model(
         detector=detector,
         gnn=gnn,
@@ -93,9 +86,9 @@ def main():
 
     trainer = Trainer(
         training_dataloader=training_dataloader,
-        validation_dataloader=validation_dataloader, 
+        validation_dataloader=validation_dataloader,
         optimizer=optimizer,
-        n_epochs=n_epochs, 
+        n_epochs=n_epochs,
         scheduler=scheduler,
         patience=patience,
     )
@@ -107,14 +100,14 @@ def main():
         pass
 
     predictor = Predictor(
-        dataloader=validation_dataloader, 
-        target=target, 
-        device=device, 
-        output_column_names=[target + '_pred', target + '_sigma'],
+        dataloader=validation_dataloader,
+        target=target,
+        device=device,
+        output_column_names=[target + '_pred'],
     )
     model._tasks[0].inference = True
     results = predictor(model)
-    save_results(db, 'dynedge_zenith', results,archive, model)
+    save_results(db, 'dynedge_energy', results,archive, model)
 
 # Main function call
 if __name__ == "__main__":

misc/badges/pylint.svg

@@ -29,18 +29,19 @@ def install(package):
     'torch-cluster',
     'torch-spline-conv',
     'torch-geometric==2.0.1',
+    'pytorch-lightning',
     'dill',
 ]
 
 # Ensure pytorch is already installed (see e.g. https://github.com/pyg-team/pytorch_geometric/issues/861#issuecomment-566424944)
 try:
     import torch  # pyright: reportMissingImports=false
 except ImportError:
-    install('torch>=1.9.0')
+    install('torch==1.9.0')
 
 setup(
     name='gnn_reco',
-    version='0.1.1',   
+    version='0.1.1',
     description='A common library for using graph neural networks (GNNs) in netrino telescopes.',
     url='https://github.com/icecube/gnn-reco',
     author='The IceCube Collaboration',

setup.py

@@ -22,11 +22,19 @@ class LossFunction(_WeightedLoss):
     """
     def __init__(
         self,
-        transform_output: Optional[Callable] = None,
+        transform_prediction_and_target: Optional[Callable] = None,
+        transform_target: Optional[Callable] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self._transform_output = transform_output if transform_output else lambda x: x
+
+        # Check(s)
+        assert not((transform_prediction_and_target is not None) and (transform_target is not None)), \
+            "Please specify at most one of `transform_prediction_and_target` and `transform_target`"
+
+        # Member variables
+        self._transform_prediction = transform_prediction_and_target if transform_prediction_and_target else lambda x: x
+        self._transform_target = transform_target if transform_target else self._transform_prediction
 
     @final
     def forward(
@@ -47,8 +55,8 @@ def forward(
             Tensor: Loss, either averaged to a scalar (if `return_elements = False`)
                 or elementwise terms with shape [N,] (if `return_elements = True`).
         """
-        prediction = self._transform_output(prediction)
-        target = self._transform_output(target)
+        prediction = self._transform_prediction(prediction)
+        target = self._transform_target(target)
 
         elements = self._forward(prediction, target)
         assert elements.size(dim=0) == target.size(dim=0), \