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

lightly_train-0.2.3/.github/pull_request_template.md

@@ -0,0 +1,7 @@
+## What has changed and why?
+
+(Delete this: Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.)
+
+## How has it been tested?
+
+(Delete this: Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration.)

{lightly_train-0.2.2 → lightly_train-0.2.3}/.gitignore

@@ -1,5 +1,6 @@
 *.code-workspace
 lightly_train_docker_test_data
+out
 
 # Created by https://www.toptal.com/developers/gitignore/api/macos,windows,linux,python
 # Edit at https://www.toptal.com/developers/gitignore?templates=macos,windows,linux,python

{lightly_train-0.2.2 → lightly_train-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lightly_train
-Version: 0.2.2
+Version: 0.2.3
 Summary: Train models with self-supervised learning in a single command
 Author: Lightly Team
 License: AGPL-3.0

{lightly_train-0.2.2 → lightly_train-0.2.3}/docker/Dockerfile-amd64-cuda

@@ -4,10 +4,8 @@
 # Start FROM PyTorch image https://hub.docker.com/r/pytorch/pytorch
 FROM pytorch/pytorch:2.3.0-cuda11.8-cudnn8-runtime AS runtime
 
-# Set the environment variable for the volume mounting
-ENV LIGHTLY_TRAIN_IS_DOCKER="True" \
-    # Install packages into the system Python and skip creating a virtual environment.
-    UV_SYSTEM_PYTHON="true" \
+# Install packages into the system Python and skip creating a virtual environment.
+ENV UV_SYSTEM_PYTHON="true" \
     # Do not cache dependencies as they would also be saved in the docker image.
     UV_NO_CACHE="true"
 
@@ -19,8 +17,8 @@ WORKDIR /home/lightly_train
 
 # Install uv
 COPY Makefile /home/lightly_train
-
 RUN make install-uv
+
 # Add uv to PATH
 ENV PATH="/root/.cargo/bin:$PATH"
 
@@ -31,8 +29,9 @@ RUN make install-docker-dependencies
 # Copy the package itself
 COPY src /home/lightly_train/src
 
+# Set and create the directory to save pretrained torch models into
+ENV TORCH_HOME="/home/lightly_train/.cache/torch"
+RUN mkdir -p ${TORCH_HOME} && chmod -R a+w $TORCH_HOME
+
 # Install the package.
 RUN make install-docker
-
-# Default command that runs when the container starts.
-ENTRYPOINT ["lightly-train"]

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

@@ -62,15 +62,26 @@ test:
 	@echo "Generate images"
 	mkdir -p $(LIGHTLY_TRAIN_DATA)
 	python -c 'from PIL import Image; [Image.new("RGB", (250, 300)).save(f"$(LIGHTLY_TRAIN_DATA)/{i}.png") for i in range(5)]'
+	@echo "Create output directory"
+	mkdir -p $(LIGHTLY_TRAIN_OUT)
+	chmod -R +rw $(LIGHTLY_TRAIN_OUT)
+	docker run --rm --shm-size=1g --user $(shell id -u):$(shell id -g) \
+		-v $(LIGHTLY_TRAIN_OUT):/out \
+		-v $(LIGHTLY_TRAIN_DATA):/data \
+		-v ./Makefile:/home/lightly_train/docker/Makefile \
+		lightly/$(IMAGE):$(TAG) make test-cli-from-within-docker -C docker
+
+# This target is run from within the docker container
+test-cli-from-within-docker:
 	@echo "Test train"
-	docker run --rm --shm-size=1g --user $(id -u):$(id -g) -v $(LIGHTLY_TRAIN_OUT):/out -v $(LIGHTLY_TRAIN_DATA):/data lightly/$(IMAGE):$(TAG) train model="torchvision/convnext_small" epochs=2 batch_size=2
-	test -f $(LIGHTLY_TRAIN_OUT)/checkpoints/last.ckpt
+	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
 	@echo "Test embed"
-	docker run --rm --shm-size=1g --user $(id -u):$(id -g) -v $(LIGHTLY_TRAIN_OUT):/out -v $(LIGHTLY_TRAIN_DATA):/data lightly/train:$(TAG) embed out="/out/embeddings.csv" checkpoint="/out/checkpoints/last.ckpt" batch_size=2 format="csv"
-	test `wc -l < $(LIGHTLY_TRAIN_OUT)/embeddings.csv` -eq 6
+	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
 	@echo "Test export"
-	docker run --rm --shm-size=1g --user $(id -u):$(id -g) -v $(LIGHTLY_TRAIN_OUT):/out -v $(LIGHTLY_TRAIN_DATA):/data lightly/train:$(TAG) export out="/out/model.pth" checkpoint="/out/checkpoints/last.ckpt" part="model" format="torch_state_dict"
-	test -f $(LIGHTLY_TRAIN_OUT)/model.pth
+	lightly-train export out="/out/model.pth" checkpoint="/out/checkpoints/last.ckpt" part="model" format="torch_state_dict"
+	test -f /out/model.pth
 
 
 

{lightly_train-0.2.2 → lightly_train-0.2.3}/docker/README.md

@@ -21,7 +21,20 @@ TODO
 
 ## Usage
 
-TODO
+First, start the docker container in interactive mode by using the -it flag. Furthermore,
+you must mount the directories you want to use.
+
+```
+docker run -it --gpus=all --user $(id -u):$(id -g) -v /my_data_dir:/data -v /my_output_dir:/out lightly/train:latest
+```
+
+Then all the usual CLI commands are fully available. E.g. run 
+
+```
+lightly-train train data="/data" out="/out" model="torchvision/convnext_small" method=dino
+```
+
+
 
 
 ## Development

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

@@ -89,3 +89,7 @@ known-first-party = ["lightly_train"]
 [tool.ruff.lint.pydocstyle]
 # Use Google-style docstrings.
 convention = "google"
+
+[tool.ruff.lint.per-file-ignores]
+# Ignore `E402` (import violations) in root `__init__.py`.
+"src/lightly_train/__init__.py" = ["E402"]

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

@@ -5,6 +5,14 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 #
+
+# Disable beta transforms warning by torchvision.
+# See https://stackoverflow.com/questions/77279407
+# TODO(Philipp, 09/24): Remove this once the warning is removed.
+import torchvision
+
+torchvision.disable_beta_transforms_warning()
+
 from lightly_train._commands.embed import embed
 from lightly_train._commands.export import ModelFormat, ModelPart, export
 from lightly_train._commands.train import train
@@ -23,4 +31,4 @@ __all__ = [
     "train",
 ]
 
-__version__ = "0.2.2"
+__version__ = "0.2.3"

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

@@ -7,7 +7,7 @@
 #
 from __future__ import annotations
 
-import warnings
+import logging
 from datetime import timedelta
 from typing import Any
 
@@ -21,6 +21,8 @@ from lightly_train._checkpoint import (
 )
 from lightly_train.types import PathLike
 
+logger = logging.getLogger(__name__)
+
 
 class ModelCheckpoint(_ModelCheckpoint):
     def __init__(
@@ -78,7 +80,6 @@ class ModelCheckpoint(_ModelCheckpoint):
                 checkpoint
             ).models
         except KeyError as ex:
-            warnings.warn(
+            logger.warning(
                 f"Could not restore lightly_train models from checkpoint: {ex}"
             )
-            pass

{lightly_train-0.2.2 → lightly_train-0.2.3}/src/lightly_train/_checkpoint.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import dataclasses
+import logging
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
@@ -21,6 +22,8 @@ from torch.serialization import MAP_LOCATION
 import lightly_train
 from lightly_train._models.embedding_model import EmbeddingModel
 
+logger = logging.getLogger(__name__)
+
 CHECKPOINT_LIGHTLY_TRAIN_KEY = "lightly_train"
 
 
@@ -104,11 +107,34 @@ class Checkpoint:
 
     @staticmethod
     def from_path(
-        checkpoint: Path, map_location: MAP_LOCATION | None = "cpu"
+        checkpoint: Path,
+        map_location: MAP_LOCATION | None = "cpu",
+        weights_only: bool = False,
     ) -> Checkpoint:
-        checkpoint_dict = torch.load(checkpoint, map_location=map_location)
+        """Load a checkpoint from a file path.
+
+        Args:
+            checkpoint:
+                Path to the checkpoint file.
+            map_location:
+                If map_location is a string, it must be a key in torch.device, such as
+                'cpu' or 'cuda:0'. If map_location is a torch.device, it will be used to
+                determine where the checkpoint should be loaded to. Default: 'cpu'.
+            weights_only:
+                If False (default), the whole checkpoint is loaded. If True, only the weights
+                of the model are loaded. This requires the user to add safe globals with:
+                https://pytorch.org/docs/stable/notes/serialization.html#torch.serialization.add_safe_globals
+                TODO(Philipp, 09/24): Expose weights_only argument to the user.
+        """
+        logger.debug(
+            f"Loading checkpoint from '{checkpoint}' with map_location '{map_location}' and weights_only {weights_only}"
+        )
+        checkpoint_dict = torch.load(
+            checkpoint, map_location=map_location, weights_only=weights_only
+        )
         return Checkpoint.from_dict(checkpoint=checkpoint_dict)
 
     def save(self, path: Path) -> None:
+        logger.debug(f"Saving checkpoint to '{path}'")
         path.parent.mkdir(parents=True, exist_ok=True)
         torch.save(self.to_dict(), path)

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

@@ -6,17 +6,23 @@
 # LICENSE file in the root directory of this source tree.
 #
 import inspect
+import logging
+import os
 import sys
 from typing import Callable
 
 from omegaconf import DictConfig, OmegaConf
 
 import lightly_train
+from lightly_train import _logging
 from lightly_train._commands import embed, export, extract_video_frames, train
 from lightly_train._commands.train import TrainConfig
+from lightly_train._logging import LIGHTLY_TRAIN_LOG_LEVEL_ENV_VAR
 from lightly_train._models import package_helpers
 from lightly_train.errors import ConfigError
 
+logger = logging.getLogger(__name__)
+
 _HELP_COMMANDS = {"help", "--help", "-h"}
 _HELP_MSG = """
     Commands:
@@ -29,6 +35,9 @@ _HELP_MSG = """
         lightly-train help                  Show help message.
 
     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.
     """
 
 _train_cfg = TrainConfig()
@@ -115,6 +124,9 @@ _TRAIN_HELP_MSG = f"""
             Additional arguments for the model. The available arguments depend on the
             the `model` parameter.
 
+    Optional arguments:
+        -v, --verbose  Run the command in verbose mode for detailed output.
+
     Examples:
     # Train a ResNet-18 model with SimCLR on ImageNet
     lightly-train train out=out data=imagenet/train model=torchvision/resnet18 method=simclr
@@ -153,6 +165,9 @@ _EXPORT_HELP_MSG = """
             and can be used to load the model with different lightly_train versions but
             requires the model to already be instantiated.
 
+    Optional arguments:
+        -v, --verbose  Run the command in verbose mode for detailed output.
+
     Examples:
     # Export the state dict of the model
     lightly-train export checkpoint=out/checkpoints/last.ckpt out=out/model.pth \\
@@ -201,6 +216,9 @@ _EMBED_HELP_MSG = """
         overwrite (bool):
             Overwrite the output file if it already exists.
 
+    Optional arguments:
+        -v, --verbose  Run the command in verbose mode for detailed output.
+
     Examples:
     # Embed images from a model checkpoint
     lightly-train embed out=embeddings.csv data=images checkpoint=out/checkpoints/last.ckpt \\
@@ -245,6 +263,9 @@ _EXTRACT_VIDEO_FRAMES_HELP_MSG = f"""
             Number of parallel calls to ffmpeg when processing multiple videos.
             If None, the number of workers is set to the number of available CPU cores.
 
+    Optional arguments:
+        -v, --verbose  Run the command in verbose mode for detailed output.
+
     Examples:
     # Extract frames from videos
     lightly-train extract_video_frames data=videos out=frames
@@ -257,12 +278,25 @@ _EXTRACT_VIDEO_FRAMES_HELP_MSG = f"""
 """
 
 
+_VERBOSE_FLAGS = ["-v", "--verbose"]
+
+
 def cli(config: DictConfig) -> None:
+    keys = list(config.keys())
+
+    # Check if the user wants to run the command in verbose mode.
+    # Any of the following will enable verbose mode: -v, --verbose
+    if any(flag in keys for flag in _VERBOSE_FLAGS):
+        os.environ[LIGHTLY_TRAIN_LOG_LEVEL_ENV_VAR] = str(logging.DEBUG)
+        config = OmegaConf.create(
+            {k: v for k, v in config.items() if k not in _VERBOSE_FLAGS}
+        )
+    _logging.set_up_console_logging()
+
     if config.is_empty():
         _show_help()
         return
 
-    keys = list(config.keys())
     # First argument after lightly_train is the command. For example `lightly-train train ...`
     command = str(keys[0]).lower()
     help_if_config_empty = True
@@ -336,17 +370,21 @@ def _run_command_fn(
     try:
         command_fn(config)
     except ConfigError as ex:
+        logger.error(ex)
+        raise ex from None  # Shorten stacktrace
+    except Exception as ex:
+        logger.error(ex)
         raise ex from None  # Shorten stacktrace
 
 
 def _list_models(config: DictConfig) -> None:
     lines = [f"    {model}" for model in package_helpers.list_model_names()]
-    print("\n".join(lines))
+    logger.info("\n".join(lines))
 
 
 def _list_methods(config: DictConfig) -> None:
     lines = [f"    {method}" for method in lightly_train.list_methods()]
-    print("\n".join(lines))
+    logger.info("\n".join(lines))
 
 
 def _is_help_command_in_config(config: DictConfig) -> bool:
@@ -370,7 +408,7 @@ def _show_invalid_command_help(command: str) -> None:
 
 
 def _show_msg(msg: str) -> None:
-    print(_format_msg(msg))
+    logger.info(_format_msg(msg))
 
 
 def _format_msg(msg: str) -> str:

{lightly_train-0.2.2 → lightly_train-0.2.3}/src/lightly_train/_commands/_warnings.py

@@ -37,6 +37,29 @@ def filter_train_warnings() -> None:
             "in this directory can be modified when the new ones are saved!"
         ),
     )
+    # Ignore mixed precision CUDA warnings as the information that CUDA is not available
+    # can be found elsewhere. The same warnings don't pop up for full precision training.
+    warnings.filterwarnings(
+        "ignore",
+        message="User provided device_type of 'cuda', but CUDA is not available.",
+        module="torch.amp",
+        category=UserWarning,
+    )
+    warnings.filterwarnings(
+        "ignore",
+        message="torch.cuda.amp.GradScaler is enabled, but CUDA is not available.",
+        module="torch.amp",
+        category=UserWarning,
+    )
+    # Ignore `lr_scheduler.step()` before `optimizer.step()` warning as it's a PyTorch Lightning issue.
+    # See https://github.com/Lightning-AI/pytorch-lightning/issues/5558
+    # TODO(Philipp, 09/24): Remove this once the issue is resolved.
+    warnings.filterwarnings(
+        "ignore",
+        message="Detected call of \`lr_scheduler.step\(\)\` before \`optimizer.step\(\)\`",
+        category=UserWarning,
+        module="torch.optim.lr_scheduler",
+    )
 
 
 def filter_embed_warnings() -> None:

{lightly_train-0.2.2 → lightly_train-0.2.3}/src/lightly_train/_commands/common_helpers.py

@@ -7,20 +7,25 @@
 #
 from __future__ import annotations
 
-import os
+import logging
+import pprint
 from pathlib import Path
+from typing import Any
 
-from omegaconf import MISSING
 from pytorch_lightning.accelerators.accelerator import Accelerator
 from pytorch_lightning.accelerators.cpu import CPUAccelerator
 from pytorch_lightning.accelerators.cuda import CUDAAccelerator
 from pytorch_lightning.accelerators.mps import MPSAccelerator
+from torch.nn import Module
 
 from lightly_train.types import PathLike
 
+logger = logging.getLogger(__name__)
+
 
 def get_checkpoint_path(checkpoint: PathLike) -> Path:
     checkpoint_path = Path(checkpoint).resolve()
+    logger.debug(f"Making sure checkpoint '{checkpoint_path}' exists.")
     if not checkpoint_path.exists():
         raise FileNotFoundError(f"Checkpoint '{checkpoint_path}' does not exist!")
     if not checkpoint_path.is_file():
@@ -30,6 +35,7 @@ def get_checkpoint_path(checkpoint: PathLike) -> Path:
 
 def get_out_path(out: PathLike, overwrite: bool) -> Path:
     out_path = Path(out).resolve()
+    logger.debug(f"Checking if output path '{out_path}' exists.")
     if out_path.exists():
         if not overwrite:
             raise ValueError(
@@ -44,37 +50,37 @@ def get_out_path(out: PathLike, overwrite: bool) -> Path:
 def get_accelerator(
     accelerator: str | Accelerator,
 ) -> str | Accelerator:
+    logger.debug(f"Getting accelerator for '{accelerator}'.")
     if accelerator != "auto":
         # User specified an accelerator, return it.
         return accelerator
 
     # Default to CUDA if available.
     if CUDAAccelerator.is_available():
+        logger.debug("CUDA is available, defaulting to CUDA.")
         return CUDAAccelerator()
     elif MPSAccelerator.is_available():
+        logger.debug("MPS is available, defaulting to MPS.")
         return MPSAccelerator()
     else:
+        logger.debug("CUDA and MPS are not available, defaulting to CPU.")
         return CPUAccelerator()
 
 
-def get_default_out() -> str:
-    return (
-        "/out" if os.getenv("LIGHTLY_TRAIN_IS_DOCKER", "False") == "True" else MISSING
-    )
-
-
-def get_default_data() -> str:
-    return (
-        "/data" if os.getenv("LIGHTLY_TRAIN_IS_DOCKER", "False") == "True" else MISSING
-    )
-
-
 def get_out_dir(out: PathLike, resume: bool, overwrite: bool) -> Path:
     out_dir = Path(out).resolve()
+    logger.debug(f"Checking if output directory '{out_dir}' exists.")
     if out_dir.exists():
         if not out_dir.is_dir():
             raise ValueError(f"Output '{out_dir}' is not a directory!")
-        dir_not_empty = any(out_dir.iterdir())
+
+        # Ignore the train.log file as it can already exist when using multiple devices.
+        # TODO(Guarin, 09/24): Fix this by checking that the directory is completely
+        # empty at the beginning. For this we have to take multiple devices and repeat
+        # calls to this function into account.
+        dir_not_empty = any(
+            filepath for filepath in out_dir.iterdir() if filepath.name != "train.log"
+        )
         if dir_not_empty and not (resume or overwrite):
             raise ValueError(
                 f"Output '{out_dir}' is not empty! Set overwrite=True to overwrite the "
@@ -82,3 +88,12 @@ def get_out_dir(out: PathLike, resume: bool, overwrite: bool) -> Path:
             )
     out_dir.mkdir(parents=True, exist_ok=True)
     return out_dir
+
+
+def pretty_format_args(
+    args: dict[str, Any], indent: int = 2, width: int = 200, compact: bool = True
+) -> str:
+    if isinstance(args.get("model"), Module):
+        args["model"] = args["model"].__class__.__name__
+
+    return f"Args: {pprint.pformat(args, indent=indent, width=width, compact=compact)}"

{lightly_train-0.2.2 → lightly_train-0.2.3}/src/lightly_train/_commands/embed.py

@@ -7,7 +7,7 @@
 #
 from __future__ import annotations
 
-import warnings
+import logging
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -21,6 +21,7 @@ from pytorch_lightning import Trainer
 from pytorch_lightning.accelerators.accelerator import Accelerator
 from torch.utils.data import DataLoader, Dataset
 
+from lightly_train import _logging
 from lightly_train._checkpoint import Checkpoint
 from lightly_train._commands import _warnings, common_helpers
 from lightly_train._commands.common_helpers import get_checkpoint_path, get_out_path
@@ -35,6 +36,8 @@ from lightly_train._embedding.writers.embedding_writer import EmbeddingWriter
 from lightly_train._models.embedding_model import EmbeddingModel
 from lightly_train.types import PathLike
 
+logger = logging.getLogger(__name__)
+
 
 def embed(
     out: PathLike,
@@ -82,7 +85,12 @@ def embed(
         overwrite:
             Overwrite the output file if it already exists.
     """
+    # Set up logging.
     _warnings.filter_embed_warnings()
+    _logging.set_up_console_logging()
+    logger.info(common_helpers.pretty_format_args(args=locals()))
+
+    logger.info(f"Embedding images in '{data}'.")
     format = _get_format(format=format)
     out_path = get_out_path(out=out, overwrite=overwrite)
     checkpoint_path = get_checkpoint_path(checkpoint=checkpoint)
@@ -103,9 +111,11 @@ def embed(
         dataloaders=dataloader,
         return_predictions=False,
     )
+    logger.info(f"Embeddings saved to '{out_path}'.")
 
 
 def embed_from_config(config: DictConfig) -> None:
+    logger.debug(f"Embedding images with config: {config}")
     config = _parse_config(config=config)
     config = _validate_config(config=config)
     config_dict = _config_to_dict(config=config)
@@ -114,8 +124,8 @@ def embed_from_config(config: DictConfig) -> None:
 
 @dataclass
 class EmbedConfig(Config):
-    out: str = common_helpers.get_default_out()
-    data: str = common_helpers.get_default_data()
+    out: str = MISSING
+    data: str = MISSING
     checkpoint: str = MISSING
     format: str = MISSING
     # OmegaConf doesn't support Union[int, Tuple[int, int]] so we have to use the
@@ -128,6 +138,7 @@ class EmbedConfig(Config):
 
 
 def _get_format(format: EmbeddingFormat | str) -> EmbeddingFormat:
+    logger.debug(f"Getting embedding format for '{format}'.")
     try:
         return EmbeddingFormat(format)
     except ValueError:
@@ -140,8 +151,10 @@ def _get_format(format: EmbeddingFormat | str) -> EmbeddingFormat:
 def _get_transform(
     image_size: int | tuple[int, int],
 ) -> EmbeddingTransform:
+    logger.debug(f"Getting embedding transform for image size {image_size}.")
     mean = tuple(IMAGENET_NORMALIZE["mean"])
     std = tuple(IMAGENET_NORMALIZE["std"])
+    logger.debug(f"Using mean {mean} and std {std} for normalization.")
     assert len(mean) == len(std) == 3
     if isinstance(image_size, int):
         image_size = (image_size, image_size)
@@ -159,8 +172,10 @@ def _get_dataset(
     transform: EmbeddingTransform,
 ) -> Dataset:
     if isinstance(data, Dataset):
+        logger.debug("Using provided dataset.")
         return data
     else:
+        logger.debug(f"Loading LightlyDataset from '{data}'.")
         return LightlyDataset(input_dir=str(data), transform=transform)
 
 
@@ -169,12 +184,15 @@ def _get_dataloader(
     batch_size: int,
     num_workers: int,
 ) -> DataLoader:
+    logger.debug(
+        f"Getting dataloader with batch_size {batch_size} and num_workers {num_workers}."
+    )
     if isinstance(dataset, Sized):
         dataset_size = len(dataset)
         if batch_size > dataset_size:
             old_batch_size = batch_size
             batch_size = dataset_size
-            warnings.warn(
+            logger.warning(
                 f"Detected dataset size {dataset_size} and batch size "
                 f"{old_batch_size}. Reducing batch size to {batch_size}."
             )
@@ -190,11 +208,13 @@ def _get_dataloader(
 
 
 def _get_embedding_model(checkpoint_path: Path) -> EmbeddingModel:
+    logger.debug(f"Loading embedding model from '{checkpoint_path}'.")
     checkpoint = Checkpoint.from_path(checkpoint=checkpoint_path)
     return checkpoint.lightly_train.models.embedding_model
 
 
 def _get_trainer(accelerator: str | Accelerator, writer: EmbeddingWriter) -> Trainer:
+    logger.debug(f"Getting trainer with accelerator '{accelerator}'.")
     return Trainer(
         accelerator=accelerator,
         devices=1,