lightly-train 0.2.3__tar.gz → 0.2.4__tar.gz

{lightly_train-0.2.3 → lightly_train-0.2.4}/.github/workflows/build_docker_image.yml

@@ -52,7 +52,7 @@ jobs:
         if: ${{ needs.build-docker-image.result == 'success' }}
         uses: rtCamp/action-slack-notify@v2
         env:
-          SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK_RELEASES }}
+          SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK_RELEASES_DEV }}
           SLACK_ICON_EMOJI: ":train:"
           SLACK_USERNAME: "Build of LightlyTrain Docker Image"
           SLACK_COLOR: "good"
@@ -71,7 +71,7 @@ jobs:
         if: ${{ needs.build-docker-image.result != 'success' }}
         uses: rtCamp/action-slack-notify@v2
         env:
-          SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK_RELEASES }}
+          SLACK_WEBHOOK: ${{ secrets.SLACK_WEBHOOK_RELEASES_DEV }}
           SLACK_ICON_EMOJI: ":x:"
           SLACK_USERNAME: "Build of LightlyTrain Docker Image"
           SLACK_COLOR: "danger"

{lightly_train-0.2.3 → lightly_train-0.2.4}/DOCS.md

@@ -27,12 +27,20 @@ lightly_train.train(
 
 In most cases you only have to specify `out`, `data`, and `model`. The rest is optional.
 
-The training process can be monitored with TensorBoard (requires 
-`pip install lightly-train[tensorboard]`):
+You can monitor your training process with the help of `tensorboard` and `wandb` loggers:
 ```
-tensorboard --logdir my_output_dir
+pip install "lightly-train[tensorboard, wandb]"
 ```
 
+Configure the loggers from Python:
+* `loggers={"tensorboard": True}`: Enable `tensorboard` logger with default arguments.
+* `loggers={"wandb": True}`: Enable `wandb` logger with default arguments.
+* `loggers={"wandb": {"project": "my-project"}}`: Configure `wandb` logger with custom arguments.
+
+LightlyTrain uses the PyTorchLightning loggers under the hood. Learn more about their configuration:
+* https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.loggers.wandb.html
+* https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.loggers.tensorboard.html
+
 #### Exporting
 ```python
 import lightly_train
@@ -94,8 +102,21 @@ lightly-train train \
 
 In most cases you only have to specify `out`, `data`, and `model`. The rest is optional.
 
-The training process can be monitored with TensorBoard (requires 
-`pip install lightly-train[tensorboard]`):
+You can monitor your training process with the help of `tensorboard` and `wandb` loggers:
+```
+pip install "lightly-train[tensorboard, wandb]"
+```
+
+Configure the loggers from the command-line:
+* `loggers.tensorboard=True`: Enable `tensorboard` logger with default arguments.
+* `loggers.wandb=True`: Enable `wandb` logger with default arguments.
+* `loggers.wandb.project="my-project"`: Configure `wandb` logger with custom arguments.
+
+LightlyTrain uses the PyTorchLightning loggers under the hood. Learn more about their configuration:
+* https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.loggers.wandb.html
+* https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.loggers.tensorboard.html
+
+
 ```
 tensorboard --logdir my_output_dir
 ```

{lightly_train-0.2.3 → lightly_train-0.2.4}/Makefile

@@ -72,7 +72,7 @@ test:
 
 .PHONY: test-ci
 test-ci:
-	pytest tests --capture=no -v
+	pytest tests -v
 
 
 ### Virtual Environment
@@ -105,8 +105,8 @@ endif
 
 # SuperGradients is not compatible with Python>=3.10. It is also not easy to install
 # on MacOS. Therefore we exclude it from the default extras.
-EXTRAS = [dev,tensorboard,timm]
-DOCKER_EXTRAS = --extra tensorboard --extra timm 
+EXTRAS = [dev,tensorboard,timm,wandb]
+DOCKER_EXTRAS = --extra tensorboard --extra timm --extra wandb
 
 # Date until which dependencies installed with --exclude-newer must have been released.
 # Dependencies released after this date are ignored.

{lightly_train-0.2.3/src/lightly_train.egg-info → lightly_train-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightly_train
-Version: 0.2.3
+Version: 0.2.4
 Summary: Train models with self-supervised learning in a single command
 Author: Lightly Team
 License: AGPL-3.0
@@ -28,6 +28,8 @@ Provides-Extra: tensorboard
 Requires-Dist: tensorboard>=2.10.0; extra == "tensorboard"
 Provides-Extra: timm
 Requires-Dist: timm>=1.0.3; extra == "timm"
+Provides-Extra: wandb
+Requires-Dist: wandb>=0.12.10; extra == "wandb"
 
 # LightlyTrain
 

{lightly_train-0.2.3 → lightly_train-0.2.4}/docker/Makefile

@@ -69,6 +69,7 @@ test:
 		-v $(LIGHTLY_TRAIN_OUT):/out \
 		-v $(LIGHTLY_TRAIN_DATA):/data \
 		-v ./Makefile:/home/lightly_train/docker/Makefile \
+		--gpus all \
 		lightly/$(IMAGE):$(TAG) make test-cli-from-within-docker -C docker
 
 # This target is run from within the docker container
@@ -76,6 +77,7 @@ test-cli-from-within-docker:
 	@echo "Test train"
 	lightly-train train data=/data out=/out model="torchvision/convnext_small" epochs=2 batch_size=2 model_args.weights="IMAGENET1K_V1" devices=2
 	test -f /out/checkpoints/last.ckpt
+	test `grep -c "GPU available: True (cuda), used: True" /out/train.log` -gt 0
 	@echo "Test embed"
 	lightly-train embed data=/data out="/out/embeddings.csv" checkpoint="/out/checkpoints/last.ckpt" batch_size=2 format="csv"
 	test `wc -l < /out/embeddings.csv` -eq 6

{lightly_train-0.2.3 → lightly_train-0.2.4}/pyproject.toml

@@ -52,6 +52,9 @@ tensorboard = [
 timm = [
     "timm>=1.0.3",
 ]
+wandb = [
+    "wandb>=0.12.10", # required by pytorch-lightning
+]
 
 [project.scripts]
 lightly-train = "lightly_train._cli:_cli_entrypoint"

{lightly_train-0.2.3 → lightly_train-0.2.4}/src/lightly_train/__init__.py

@@ -31,4 +31,4 @@ __all__ = [
     "train",
 ]
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"

lightly_train-0.2.4/src/lightly_train/_callbacks/callback_args.py

@@ -0,0 +1,45 @@
+#
+# Copyright (c) Lightly AG and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+from dataclasses import dataclass, field
+from typing import Optional
+
+from lightly_train._callbacks.checkpoint import ModelCheckpointArgs
+from lightly_train._configs.config import Config
+
+
+@dataclass
+class LearningRateMonitorArgs(Config):
+    pass
+
+
+@dataclass
+class DeviceStatsMonitorArgs(Config):
+    pass
+
+
+@dataclass
+class EarlyStoppingArgs(Config):
+    monitor: str = "train_loss"
+    patience: int = int(1e12)
+    check_finite: bool = True
+
+
+@dataclass
+class CallbackArgs(Config):
+    learning_rate_monitor: Optional[LearningRateMonitorArgs] = field(
+        default_factory=LearningRateMonitorArgs
+    )
+    device_stats_monitor: Optional[DeviceStatsMonitorArgs] = field(
+        default_factory=DeviceStatsMonitorArgs
+    )
+    early_stopping: Optional[EarlyStoppingArgs] = field(
+        default_factory=EarlyStoppingArgs
+    )
+    model_checkpoint: Optional[ModelCheckpointArgs] = field(
+        default_factory=ModelCheckpointArgs
+    )

lightly_train-0.2.4/src/lightly_train/_callbacks/callback_helpers.py

@@ -0,0 +1,75 @@
+#
+# Copyright (c) Lightly AG and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+from __future__ import annotations
+
+import dataclasses
+from pathlib import Path
+from typing import Any
+
+from omegaconf import OmegaConf
+from pytorch_lightning import Callback
+from pytorch_lightning.callbacks import (
+    DeviceStatsMonitor,
+    EarlyStopping,
+    LearningRateMonitor,
+)
+from torch.nn import Module
+
+from lightly_train._callbacks.callback_args import (
+    CallbackArgs,
+)
+from lightly_train._callbacks.checkpoint import ModelCheckpoint
+from lightly_train._checkpoint import CheckpointLightlyTrainModels
+from lightly_train._configs import validate
+from lightly_train._models.embedding_model import EmbeddingModel
+
+
+def validate_callback_args(callback_dict: dict[str, Any] | None) -> CallbackArgs:
+    callback_dictconfig = OmegaConf.create(
+        {} if callback_dict is None else callback_dict
+    )
+    callback_dictconfig = validate.validate_dictconfig(
+        callback_dictconfig, CallbackArgs
+    )
+    callback_args = OmegaConf.to_object(callback_dictconfig)
+    assert isinstance(callback_args, CallbackArgs)  # make mypy happy
+    return callback_args
+
+
+def get_callbacks(
+    callback_args: CallbackArgs,
+    out: Path,
+    model: Module,
+    embedding_model: EmbeddingModel,
+):
+    callbacks: list[Callback] = []
+    if callback_args.learning_rate_monitor is not None:
+        callbacks.append(
+            LearningRateMonitor(
+                **dataclasses.asdict(callback_args.learning_rate_monitor)
+            )
+        )
+    if callback_args.device_stats_monitor is not None:
+        callbacks.append(
+            DeviceStatsMonitor(**dataclasses.asdict(callback_args.device_stats_monitor))
+        )
+    if callback_args.early_stopping is not None:
+        callbacks.append(
+            EarlyStopping(**dataclasses.asdict(callback_args.early_stopping))
+        )
+    if callback_args.model_checkpoint is not None:
+        callbacks.append(
+            ModelCheckpoint(
+                models=CheckpointLightlyTrainModels(
+                    model=model, embedding_model=embedding_model
+                ),
+                dirpath=out / "checkpoints",
+                **dataclasses.asdict(callback_args.model_checkpoint),
+            )
+        )
+    return callbacks

{lightly_train-0.2.3 → lightly_train-0.2.4}/src/lightly_train/_callbacks/checkpoint.py

@@ -8,8 +8,9 @@
 from __future__ import annotations
 
 import logging
+from dataclasses import dataclass
 from datetime import timedelta
-from typing import Any
+from typing import Any, Optional
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint as _ModelCheckpoint
@@ -19,11 +20,20 @@ from lightly_train._checkpoint import (
     CheckpointLightlyTrain,
     CheckpointLightlyTrainModels,
 )
+from lightly_train._configs.config import Config
 from lightly_train.types import PathLike
 
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class ModelCheckpointArgs(Config):
+    save_last: bool = True
+    enable_version_counter: bool = False
+    save_top_k: int = 1
+    every_n_epochs: Optional[int] = None
+
+
 class ModelCheckpoint(_ModelCheckpoint):
     def __init__(
         self,

{lightly_train-0.2.3 → lightly_train-0.2.4}/src/lightly_train/_cli.py

@@ -34,7 +34,7 @@ _HELP_MSG = """
         lightly-train extract_video_frames  Extract frames from videos using ffmpeg.
         lightly-train help                  Show help message.
 
-    Run `lightly_train <command> help` for more information on a specific command.
+    Run `lightly-train <command> help` for more information on a specific command.
 
     Optional arguments:
         -v, --verbose  Run the command in verbose mode for detailed output.
@@ -105,6 +105,20 @@ _TRAIN_HELP_MSG = f"""
             For details, see: https://lightning.ai/docs/pytorch/stable/common/trainer.html#precision
         seed (int):
             Random seed for reproducibility. Default: {_train_cfg.seed}
+        loggers (dict):
+            Loggers for training. Either null or a dictionary of logger names to either
+            null or a dictionary of logger arguments. Null uses the default loggers.
+            To disable a logger, set it to null: `loggers.tensorboard=null`.
+            To configure a logger, pass the respective arguments:
+            `loggers.wandb.project="my_project"`.
+            Default: null
+        callbacks (dict):
+            Callbacks fo training. Either null or a dictionary of callback names to
+            either null or a dictionary of callback arguments. Null uses the default
+            callbacks. To disable a callback, set it to null:
+            `callbacks.model_checkpoint=null`. To configure a callback, pass the
+            respective arguments: `callbacks.model_checkpoint.every_n_epochs=5`.
+            Default: null
         optim_args (dict):
             Arguments for AdamW optimizer. Available arguments are:
             - lr (float)
@@ -159,10 +173,10 @@ _EXPORT_HELP_MSG = """
             Format to save the model in. Valid options are 'torch_model' and
             'torch_state_dict'. 'torch_model' saves the model as a torch module which
             can be loaded with `model = torch.load(out)`. This requires that the same
-            lightly_train version is installed when the model is exported and when it is
+            LightlyTrain version is installed when the model is exported and when it is
             loaded again. 'torch_state_dict' saves the model's state dict which can be
             loaded with `model.load_state_dict(torch.load(out))`. This is more flexible
-            and can be used to load the model with different lightly_train versions but
+            and can be used to load the model with different LightlyTrain versions but
             requires the model to already be instantiated.
 
     Optional arguments:
@@ -399,7 +413,7 @@ def _show_invalid_command_help(command: str) -> None:
     msg = _format_msg(
         f"""
         Unknown command '{command}':
-            lightly_train {command}
+            lightly-train {command}
         """
     )
     msg += "\n"

{lightly_train-0.2.3 → lightly_train-0.2.4}/src/lightly_train/_commands/export.py

@@ -52,10 +52,10 @@ def export(
             Format to save the model in. Valid options are 'torch_model' and
             'torch_state_dict'. 'torch_model' saves the model as a torch module which
             can be loaded with `model = torch.load(out)`. This requires that the same
-            lightly_train version is installed when the model is exported and when it is
+            LightlyTrain version is installed when the model is exported and when it is
             loaded again. 'torch_state_dict' saves the model's state dict which can be
             loaded with `model.load_state_dict(torch.load(out))`. This is more flexible
-            and can be used to load the model with different lightly_train versions but
+            and can be used to load the model with different LightlyTrain versions but
             requires the model to already be instantiated.
         overwrite:
             Overwrite the output file if it already exists.

{lightly_train-0.2.3 → lightly_train-0.2.4}/src/lightly_train/_commands/train.py

@@ -21,7 +21,9 @@ from pytorch_lightning.trainer.connectors.accelerator_connector import _PRECISIO
 from torch.nn import Module
 from torch.utils.data import Dataset
 
+import lightly_train._loggers.logger_helpers
 from lightly_train import _logging
+from lightly_train._callbacks import callback_helpers
 from lightly_train._commands import _warnings, common_helpers, train_helpers
 from lightly_train._configs import omegaconf_utils, validate
 from lightly_train._configs.config import Config
@@ -49,6 +51,8 @@ def train(
     strategy: str | Strategy = "auto",
     precision: _PRECISION_INPUT = "32-true",  # Default precision in PyTorch Lightning
     seed: int = 0,
+    loggers: dict[str, dict[str, Any] | None] | None = None,
+    callbacks: dict[str, dict[str, Any] | None] | None = None,
     optim_args: dict[str, Any] | None = None,
     transform_args: dict[str, Any] | None = None,
     loader_args: dict[str, Any] | None = None,
@@ -111,6 +115,19 @@ def train(
             https://lightning.ai/docs/pytorch/stable/common/trainer.html#precision
         seed:
             Random seed for reproducibility.
+        loggers:
+            Loggers for training. Either None or a dictionary of logger names to either
+            None or a dictionary of logger arguments. None uses the default loggers.
+            To disable a logger, set it to None: `loggers={"tensorboard": None}`.
+            To configure a logger, pass the respective arguments:
+            `loggers={"wandb": {"project": "my_project"}}`.
+        callbacks:
+            Callbacks for training. Either None or a dictionary of callback names to
+            either None or a dictionary of callback arguments. None uses the default
+            callbacks. To disable a callback, set it to None:
+            `callbacks={"model_checkpoint": None}`. To configure a callback, pass the
+            respective arguments:
+            `callbacks={"model_checkpoint": {"every_n_epochs": 5}}`.
         optim_args:
             Arguments for AdamW optimizer. Available arguments are:
             - lr: float
@@ -160,10 +177,21 @@ def train(
     log_every_n_steps = train_helpers.get_lightning_logging_interval(
         dataset_size=scaling_info.dataset_size, batch_size=batch_size
     )
-    trainer_instance = train_helpers.get_trainer(
+    logger_args = lightly_train._loggers.logger_helpers.validate_logger_args(
+        loggers=loggers
+    )
+    logger_instances = lightly_train._loggers.logger_helpers.get_loggers(
+        logger_args=logger_args, out=out_dir
+    )
+    callback_args = callback_helpers.validate_callback_args(callback_dict=callbacks)
+    callback_instances = callback_helpers.get_callbacks(
+        callback_args=callback_args,
         out=out_dir,
         model=model_instance,
         embedding_model=embedding_model,
+    )
+    trainer_instance = train_helpers.get_trainer(
+        out=out_dir,
         epochs=epochs,
         accelerator=accelerator,
         strategy=strategy,
@@ -171,6 +199,8 @@ def train(
         num_nodes=num_nodes,
         precision=precision,
         log_every_n_steps=log_every_n_steps,
+        loggers=logger_instances,
+        callbacks=callback_instances,
         trainer_args=trainer_args,
     )
     dataloader = train_helpers.get_dataloader(
@@ -228,6 +258,8 @@ class TrainConfig(Config):
     strategy: str = "auto"
     precision: str = "32-true"
     seed: int = 0
+    loggers: Optional[Dict[str, Optional[Dict[str, Any]]]] = None
+    callbacks: Optional[Dict[str, Optional[Dict[str, Any]]]] = None
     optim_args: Optional[Dict[str, Any]] = None
     transform_args: Optional[Dict[str, Any]] = None
     loader_args: Optional[Dict[str, Any]] = None

{lightly_train-0.2.3 → lightly_train-0.2.4}/src/lightly_train/_commands/train_helpers.py

@@ -7,32 +7,24 @@
 #
 from __future__ import annotations
 
+import logging
 from pathlib import Path
 from typing import Any, Callable, Sized, Type
 
 from lightly.data import LightlyDataset
-from pytorch_lightning import Trainer
+from pytorch_lightning import Callback, Trainer
 from pytorch_lightning.accelerators.accelerator import Accelerator
 from pytorch_lightning.accelerators.cpu import CPUAccelerator
 from pytorch_lightning.accelerators.cuda import CUDAAccelerator
-from pytorch_lightning.callbacks import (
-    DeviceStatsMonitor,
-    EarlyStopping,
-    LearningRateMonitor,
-)
 from pytorch_lightning.loggers import Logger
 from pytorch_lightning.strategies.strategy import Strategy
 from pytorch_lightning.trainer.connectors.accelerator_connector import _PRECISION_INPUT
 from torch.nn import Module
 from torch.utils.data import DataLoader, Dataset
 
-from lightly_train._callbacks.checkpoint import ModelCheckpoint
-from lightly_train._checkpoint import CheckpointLightlyTrainModels
 from lightly_train._commands import common_helpers
 from lightly_train._configs import validate
 from lightly_train._constants import DATALOADER_TIMEOUT
-from lightly_train._loggers.jsonl import JSONLLogger
-from lightly_train._loggers.tensorboard import TensorBoardLogger
 from lightly_train._methods import method_helpers
 from lightly_train._methods.method import Method
 from lightly_train._models import package_helpers
@@ -42,18 +34,6 @@ from lightly_train._optim.optimizer_type import OptimizerType
 from lightly_train._scaling import IMAGENET_SIZE, ScalingInfo
 from lightly_train.types import PathLike, Transform
 
-try:
-    import timm
-except ImportError:
-    timm = None
-
-try:
-    import tensorboard
-except ImportError:
-    tensorboard = None
-
-import logging
-
 logger = logging.getLogger(__name__)
 
 
@@ -166,8 +146,6 @@ def get_embedding_model(model: Module, embed_dim: int | None = None) -> Embeddin
 
 def get_trainer(
     out: Path,
-    model: Module,
-    embedding_model: EmbeddingModel,
     epochs: int,
     accelerator: str | Accelerator,
     strategy: str | Strategy,
@@ -175,17 +153,11 @@ def get_trainer(
     num_nodes: int,
     log_every_n_steps: int,
     precision: _PRECISION_INPUT | None,
+    loggers: list[Logger],
+    callbacks: list[Callback],
     trainer_args: dict[str, Any] | None,
 ) -> Trainer:
     logger.debug("Getting trainer.")
-    # Set version and name to empty string to save logs directly in the root
-    # directory.
-    loggers: list[Logger] = [
-        JSONLLogger(save_dir=out, name="", version=""),
-    ]
-    if tensorboard is not None:
-        loggers.append(TensorBoardLogger(save_dir=out, name="", version=""))
-    logger.debug(f"Using loggers {[log.__class__.__name__ for log in loggers]}.")
 
     accelerator = common_helpers.get_accelerator(accelerator=accelerator)
     strategy = get_strategy(accelerator=accelerator, strategy=strategy, devices=devices)
@@ -200,20 +172,7 @@ def get_trainer(
         num_nodes=num_nodes,
         precision=precision,
         log_every_n_steps=log_every_n_steps,
-        callbacks=[
-            LearningRateMonitor(),
-            DeviceStatsMonitor(),
-            ModelCheckpoint(
-                save_last=True,
-                models=CheckpointLightlyTrainModels(
-                    model=model, embedding_model=embedding_model
-                ),
-                dirpath=out / "checkpoints",
-                enable_version_counter=False,
-            ),
-            # Stop if training loss diverges.
-            EarlyStopping(monitor="train_loss", patience=int(1e12), check_finite=True),
-        ],
+        callbacks=callbacks,
         logger=loggers,
         sync_batchnorm=sync_batchnorm,
     )

{lightly_train-0.2.3 → lightly_train-0.2.4}/src/lightly_train/_loggers/jsonl.py

@@ -5,9 +5,12 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 #
+from __future__ import annotations
+
 import json
 import logging
 import os
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 from lightning_fabric.loggers.logger import rank_zero_experiment
@@ -15,6 +18,7 @@ from lightning_fabric.utilities.rank_zero import rank_zero_warn
 from pytorch_lightning.loggers import CSVLogger
 from pytorch_lightning.loggers.csv_logs import ExperimentWriter as CSVExperimentWriter
 
+from lightly_train._configs.config import Config
 from lightly_train.types import PathLike
 
 log = logging.getLogger(__name__)
@@ -55,6 +59,11 @@ class ExperimentWriter(CSVExperimentWriter):
             )
 
 
+@dataclass
+class JSONLLoggerArgs(Config):
+    flush_logs_every_n_steps: int = 100
+
+
 class JSONLLogger(CSVLogger):
     """Log to local file system in JSON Lines format.
 
@@ -69,7 +78,7 @@ class JSONLLogger(CSVLogger):
 
     Args:
         save_dir: Save directory
-        name: Experiment name. Defaults to ``'lightning_logs'``.
+        name: Experiment name.
         version:
             Experiment version. If version is not specified the logger inspects the save
             directory for existing versions, then automatically assigns the next
@@ -83,8 +92,8 @@ class JSONLLogger(CSVLogger):
     def __init__(
         self,
         save_dir: PathLike,
-        name: str = "lightning_logs",
-        version: Optional[Union[int, str]] = None,
+        name: str = "",
+        version: Optional[Union[int, str]] = "",
         prefix: str = "",
         flush_logs_every_n_steps: int = 100,
     ):

lightly_train-0.2.4/src/lightly_train/_loggers/logger_args.py

@@ -0,0 +1,23 @@
+#
+# Copyright (c) Lightly AG and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+from dataclasses import dataclass, field
+from typing import Optional
+
+from lightly_train._configs.config import Config
+from lightly_train._loggers.jsonl import JSONLLoggerArgs
+from lightly_train._loggers.tensorboard import TensorBoardLoggerArgs
+from lightly_train._loggers.wandb import WandbLoggerArgs
+
+
+@dataclass
+class LoggerArgs(Config):
+    jsonl: Optional[JSONLLoggerArgs] = field(default_factory=JSONLLoggerArgs)
+    tensorboard: Optional[TensorBoardLoggerArgs] = field(
+        default_factory=TensorBoardLoggerArgs
+    )
+    wandb: Optional[WandbLoggerArgs] = None