salt-ml 0.6__py3-none-any.whl

salt/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.6"

salt/main.py

@@ -0,0 +1,31 @@
+"""Top level training script, powered by the lightning CLI."""
+
+import pathlib
+
+import comet_ml  # noqa: F401
+from lightning.pytorch.cli import ArgsType
+
+from salt.callbacks import SaveConfigCallback
+from salt.data.datamodules import SaltDataModule
+from salt.modelwrapper import ModelWrapper
+from salt.utils.cli import SaltCLI
+
+
+def main(args: ArgsType = None) -> None:
+    config_dir = pathlib.Path(__file__).parent / "configs"
+
+    SaltCLI(
+        model_class=ModelWrapper,
+        datamodule_class=SaltDataModule,
+        save_config_callback=SaveConfigCallback,
+        args=args,
+        parser_kwargs={
+            "default_env": True,
+            "fit": {"default_config_files": [f"{config_dir}/base.yaml"]},
+            "test": {"default_config_files": [f"{config_dir}/base.yaml"]},
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

salt/modelwrapper.py

@@ -0,0 +1,252 @@
+import math
+import warnings
+from collections.abc import Mapping
+
+import lightning as L
+import torch
+from torch import nn
+
+from salt.models import InputNorm
+
+
+def check_unique(modules: nn.ModuleList, attr_name: str) -> None:
+    assert len({getattr(m, attr_name) for m in modules}) == len(
+        modules
+    ), f"Attribute '{attr_name}' must be unique for class {modules[0].__class__.__name__}"
+
+
+class ModelWrapper(L.LightningModule):
+    def __init__(
+        self,
+        model: nn.Module,
+        lrs_config: Mapping[str, float],
+        global_object: str,
+        norm_config: dict | None = None,
+        name: str = "salt",
+        muP_config: dict | None = None,
+        loss_mode: str = "wsum",
+    ):
+        """A wrapper class for any model implemented in Salt.
+
+        This wrapper class allows is as generic as possible. It wraps
+        [`SaltModel`][salt.models.SaltModel], but could also be used to
+        wrap any other model if you want to do train something that doesn't
+        fit into the [`SaltModel`][salt.models.SaltModel] architecture.
+
+        This class is responsible for containing things that are common to all
+        salt models. These are:
+
+        - A generic forward pass, including input normalisation
+        - Training, validation and test steps, which include logging
+        - Some sanity checks on the model configuration
+
+        Parameters
+        ----------
+        model : nn.Module
+            Model to be wrapped
+        lrs_config: Mapping
+            LRS config which has to be set manually for now
+            https://github.com/omni-us/jsonargparse/issues/170#issuecomment-1288167674
+        global_object : str
+            Name of the global input object, as opposed to the constituent-level
+            inputs. This argument is set automatically by the framework.
+        norm_config : dict, optional
+            Keyword arguments for [`salt.models.InputNorm`][salt.models.InputNorm].
+        name: str, optional
+            Name of the model, used for logging and inference output names
+        muP_config: dict, optional
+            The muP configuration.
+        loss_mode: str, optional
+            The loss mode to use. Default is "wsum" (weighted sum).
+            Other options are
+            - 'GLS' : arxiv.org/1904.08492
+        """
+        super().__init__()
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            self.save_hyperparameters(logger=False)
+
+        self.model = model
+        self.lrs_config = lrs_config
+        self.global_object = global_object
+        self.name = name
+        self.muP = muP_config or {}
+        self.last_val_batch_outs = None
+        # Here the model should pick it up
+        if self.muP:
+            from salt.utils.muP_utils.configuration_muP import instantiate_mup
+
+            load_path = None
+            if "shape_path" in self.muP:
+                load_path = self.muP["shape_path"]
+            instantiate_mup(model, load_path)
+
+        # all tasks should inherit the global object type and model name
+        self.model.global_object = self.global_object
+        for task in self.model.tasks:
+            task.global_object = self.global_object
+            task.model_name = self.name
+
+        # ensure unique names for init_nets and tasks
+        check_unique(self.model.init_nets, "input_name")
+        check_unique(self.model.tasks, "name")
+
+        # check that the model has the same output size for all init nets
+        assert len({t.net.output_size for t in self.model.init_nets if t.input_name != "EDGE"}) == 1
+
+        # create input normaliser
+        assert norm_config is not None
+        self.norm = InputNorm(**norm_config)
+        allowed_loss_modes = ["wsum", "GLS"]
+        assert loss_mode in allowed_loss_modes, f"Loss mode must be one of {allowed_loss_modes}"
+        self.loss_mode = loss_mode
+        if loss_mode == "GLS":
+            assert all(
+                task.weight == 1.0 for task in self.model.tasks
+            ), "GLS does not utilise task weights - remove all/set to 1"
+
+    def total_loss(self, loss: dict):
+        """Computes the final loss based on the loss mode."""
+        if self.loss_mode == "GLS":
+            # Calculate the geometric mean of the losses
+            loss_prod = math.prod(subloss for subloss in loss.values())
+            return torch.pow(loss_prod, 1.0 / len(loss))
+
+        # Return the default weighted sum
+        return sum(subloss for subloss in loss.values())
+
+    def forward(self, inputs, pad_masks=None, labels=None):
+        """Generic forward pass through any salt-compatible model.
+
+        This function performs input normalisation and then calls the `self.model`'s
+        forward pass. Don't call this method directy, instead use `__call__`.
+
+        Parameters
+        ----------
+        inputs
+            Any generic input to the model.
+        pad_masks
+            Input padding masks.
+        labels
+            Training targets. If not specified, assume we are running model inference
+            (i.e. no loss computation).
+
+        Returns
+        -------
+        Whatever is returned by `self.model`'s forward pass.
+        """
+        x = self.norm(inputs)
+        return self.model(x, pad_masks, labels)
+
+    def shared_step(self, batch, evaluation=False):
+        """Function used to unpack the batch, run the forward pass, and compute
+        losses, used by training, validation and test steps.
+
+        Parameters
+        ----------
+        batch : tuple
+            A single batch of inputs, pad_masks and labels
+        evaluation : bool
+            If true, don't compute the losses and return early
+
+        Returns
+        -------
+        preds
+            Model predictions
+        labels
+            True labels
+        loss
+            Reduced loss over the input batch
+        """
+        # unpack the batch
+        inputs, pad_masks, labels = batch
+
+        # forward pass through model
+        preds, loss = self(inputs, pad_masks, labels)
+
+        if evaluation:
+            return preds, labels, pad_masks, None
+
+        # compute total loss
+        loss["loss"] = self.total_loss(loss)
+
+        return preds, labels, pad_masks, loss
+
+    def log_losses(self, loss, stage):
+        kwargs = {"sync_dist": len(self.trainer.device_ids) > 1}
+        self.log(f"{stage}_loss", loss["loss"], **kwargs)
+        for t, loss_value in loss.items():
+            n = f"{stage}_{t}_loss" if "loss" not in t else f"{stage}_{t}"
+            self.log(n, loss_value, **kwargs)
+
+    def training_step(self, batch):
+        # foward pass
+        preds, labels, pad_masks, loss = self.shared_step(batch)
+
+        if loss["loss"].isnan():
+            raise RuntimeError(
+                "Loss is NaN - this indicates something significant has gone wrong."
+                "Check for any NaNs or infs in the input dataset. If nothing is found here, "
+                "check 'docs/training.md - NaNs' for more information"
+            )
+
+        # log losses
+        self.log_losses(loss, stage="train")
+
+        outputs = {
+            "preds": preds,
+            "labels": labels,
+            "pad_masks": pad_masks,
+        }
+
+        return {**loss, "outputs": outputs}
+
+    def validation_step(self, batch):
+        # foward pass
+        preds, labels, pad_masks, loss = self.shared_step(batch)
+
+        # log losses
+        self.log_losses(loss, stage="val")
+
+        # Store outputs to be used by the MaskformerMetrics callback
+        outputs = {
+            "preds": preds,
+            "labels": labels,
+            "pad_masks": pad_masks,
+        }
+
+        return {**loss, "outputs": outputs}
+
+    def test_step(self, batch):
+        inputs, pad_masks, _ = batch
+        batch = (inputs, pad_masks, None)
+        return self.shared_step(batch, evaluation=True)[0]
+
+    def configure_optimizers(self):
+        if self.muP:
+            from mup import MuAdamW as AdamW
+        else:
+            from torch.optim import AdamW
+        opt = AdamW(
+            self.parameters(),
+            lr=self.lrs_config["initial"],
+            weight_decay=self.lrs_config.get("weight_decay", 1e-5),
+        )
+
+        # 1cycle
+        sch = torch.optim.lr_scheduler.OneCycleLR(
+            opt,
+            max_lr=self.lrs_config["max"],
+            total_steps=self.trainer.estimated_stepping_batches,
+            div_factor=self.lrs_config["max"] / self.lrs_config["initial"],
+            final_div_factor=self.lrs_config["initial"] / self.lrs_config["end"],
+            pct_start=float(self.lrs_config["pct_start"]),
+            last_epoch=int(self.lrs_config.get("last_epoch", -1)),
+        )
+        sch = {"scheduler": sch, "interval": "step"}
+
+        return [opt], [sch]
+
+    @property
+    def input_dims(self) -> dict[str, int]:
+        return {k: len(v) for k, v in self.norm.variables.items()}

salt/stypes.py

@@ -0,0 +1,8 @@
+from typing import TypeAlias
+
+from torch import BoolTensor, Tensor
+
+Vars: TypeAlias = dict[str, list[str]]
+Tensors: TypeAlias = dict[str, Tensor]
+BoolTensors: TypeAlias = dict[str, BoolTensor]
+NestedTensors: TypeAlias = dict[str, dict[str, Tensor]]

salt_ml-0.6.dist-info/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2024 CERN
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

salt_ml-0.6.dist-info/METADATA

@@ -0,0 +1,32 @@
+Metadata-Version: 2.1
+Name: salt-ml
+Version: 0.6
+Summary: Multimodal and Multiclass Machine Learning for High Energy Physics
+License: MIT
+Project-URL: Homepage, https://gitlab.cern.ch/atlas-flavor-tagging-tools/algorithms/salt
+Project-URL: Issue Tracker, https://gitlab.cern.ch/atlas-flavor-tagging-tools/algorithms/salt/-/issues
+Requires-Python: <3.12,>=3.10
+License-File: LICENSE
+Requires-Dist: pytest ==7.4.2
+Requires-Dist: pytest-cov ==4.1.0
+Requires-Dist: comet-ml ==3.35.3
+Requires-Dist: tqdm ==4.66.1
+Requires-Dist: rich ==13.6.0
+Requires-Dist: pre-commit ==3.5.0
+Requires-Dist: ruff ==0.2.2
+Requires-Dist: h5py ==3.10.0
+Requires-Dist: numpy ==1.24.2
+Requires-Dist: torch ==2.2.1
+Requires-Dist: lightning ==2.2.0
+Requires-Dist: jsonargparse[all] ==4.32.0
+Requires-Dist: torchmetrics ==1.2.1
+Requires-Dist: onnx ==1.16.0
+Requires-Dist: onnxruntime ==1.15.1
+Requires-Dist: atlas-ftag-tools ==0.2.2
+Requires-Dist: scipy ==1.12.0
+Requires-Dist: mup ==1.0.0
+Requires-Dist: boto3 ==1.28.17
+Requires-Dist: s3fs ==2023.9.2
+Requires-Dist: s3path ==0.5.0
+Requires-Dist: tensorboard ==2.14.1
+

salt_ml-0.6.dist-info/RECORD

@@ -0,0 +1,10 @@
+salt/__init__.py,sha256=ncKOhmMshkKH0hADIDc1EohSZbPVhx6HkN9Q9lGRlpc,20
+salt/main.py,sha256=k080VQ4T6Qt7tFOwFnuDZeIEKf5x9Yyr9Zx7kDr99v0,849
+salt/modelwrapper.py,sha256=rxL6hk7tTiEpLKmADqfH0wC9_b-eGRs2cb3IsQxr9D8,8849
+salt/stypes.py,sha256=YrnZOfdpYNFrpanHeQEKc0oHAyiyXFMtpJr3T18INoc,249
+salt_ml-0.6.dist-info/LICENSE,sha256=xhKDC68K3tkxStAz7NqUflkVqfvXR_Zw4nsWbsisXq8,1052
+salt_ml-0.6.dist-info/METADATA,sha256=XOI1YNi7hA77jnm1jNKZvS1Q4ZBxkFARBlI-3-4BKAU,1089
+salt_ml-0.6.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+salt_ml-0.6.dist-info/entry_points.txt,sha256=U3e3HaUskBtZQMhevO_XNhXDZyxP4DNVpe9USQZIVEE,317
+salt_ml-0.6.dist-info/top_level.txt,sha256=BqSzl_xwRerFJ8Kuw_9GzuXOMAFsdg8jjCbnmQL-Z7Y,5
+salt_ml-0.6.dist-info/RECORD,,

salt_ml-0.6.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.1.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

salt_ml-0.6.dist-info/entry_points.txt

@@ -0,0 +1,8 @@
+[console_scripts]
+compare_models = salt.utils.compare_models:main
+download_S3 = salt.utils.file_utils:download_from_S3
+get_onnx_metadata = salt.utils.get_onnx_metadata:main
+repair_ckpt = salt.utils.repair_ckpt:main
+salt = salt.main:main
+setup_muP = salt.utils.muP_utils.main_muP:main
+to_onnx = salt.onnx.to_onnx:main

salt_ml-0.6.dist-info/top_level.txt

@@ -0,0 +1 @@
+salt