├── src └── lightly_train │ ├── py.typed │ ├── _models │ ├── dinov2_vit │ │ ├── dinov2_vit_src │ │ │ ├── configs │ │ │ │ ├── train │ │ │ │ │ ├── vits14.yaml │ │ │ │ │ ├── vitb14.yaml │ │ │ │ │ ├── _vittest14.yaml │ │ │ │ │ ├── vits14_reg4.yaml │ │ │ │ │ ├── vitb14_reg4.yaml │ │ │ │ │ ├── vitg14.yaml │ │ │ │ │ ├── vitl14.yaml │ │ │ │ │ ├── vitg14_reg4.yaml │ │ │ │ │ └── vitl14_reg4.yaml │ │ │ │ └── eval │ │ │ │ │ ├── vitb14_pretrain.yaml │ │ │ │ │ ├── vitl14_pretrain.yaml │ │ │ │ │ ├── vits14_pretrain.yaml │ │ │ │ │ ├── vitg14_pretrain.yaml │ │ │ │ │ ├── vitb14_reg4_pretrain.yaml │ │ │ │ │ ├── vitl14_reg4_pretrain.yaml │ │ │ │ │ ├── vits14_reg4_pretrain.yaml │ │ │ │ │ └── vitg14_reg4_pretrain.yaml │ │ │ ├── layers │ │ │ │ ├── layer_scale.py │ │ │ │ ├── __init__.py │ │ │ │ ├── drop_path.py │ │ │ │ └── mlp.py │ │ │ └── models │ │ │ │ └── __init__.py │ │ └── __init__.py │ ├── dinov3 │ │ ├── dinov3_src │ │ │ ├── __init__.py │ │ │ ├── hub │ │ │ │ ├── __init__.py │ │ │ │ └── utils.py │ │ │ ├── utils │ │ │ │ ├── __init__.py │ │ │ │ ├── dtype.py │ │ │ │ └── custom_callable.py │ │ │ └── layers │ │ │ │ ├── rms_norm.py │ │ │ │ ├── layer_scale.py │ │ │ │ └── __init__.py │ │ ├── __init__.py │ │ └── dinov3_convnext.py │ ├── __init__.py │ ├── timm │ │ └── __init__.py │ ├── custom │ │ └── __init__.py │ ├── rfdetr │ │ └── __init__.py │ ├── super_gradients │ │ ├── __init__.py │ │ └── super_gradients.py │ ├── torchvision │ │ ├── __init__.py │ │ ├── torchvision.py │ │ ├── resnet.py │ │ ├── convnext.py │ │ └── shufflenet.py │ ├── ultralytics │ │ └── __init__.py │ ├── log_usage_example.py │ ├── rtdetr │ │ └── rtdetr.py │ └── package.py │ ├── _configs │ ├── __init__.py │ └── omegaconf_utils.py │ ├── _data │ ├── __init__.py │ ├── _serialize │ │ └── __init__.py │ ├── label_helpers.py │ ├── task_data_args.py │ ├── infinite_cycle_iterator.py │ ├── cache.py │ ├── download.py │ └── task_dataset.py │ ├── _loggers │ ├── __init__.py │ ├── wandb.py │ ├── logger_args.py │ ├── tensorboard.py │ └── task_logger_args.py │ ├── _methods │ ├── __init__.py │ ├── dino │ │ └── __init__.py │ ├── densecl │ │ ├── __init__.py │ │ └── densecl_loss.py │ ├── dinov2 │ │ ├── __init__.py │ │ ├── scheduler.py │ │ └── dinov2_transform.py │ ├── simclr │ │ └── __init__.py │ ├── distillation │ │ └── __init__.py │ ├── distillationv2 │ │ ├── __init__.py │ │ └── distillationv2_loss.py │ └── method_args.py │ ├── _optim │ ├── __init__.py │ ├── optimizer_type.py │ ├── trainable_modules.py │ ├── sgd_args.py │ ├── optimizer_args.py │ ├── lars_args.py │ └── adamw_args.py │ ├── _callbacks │ ├── __init__.py │ ├── learning_rate_monitor.py │ ├── callback_args.py │ ├── export.py │ └── mlflow_logging.py │ ├── _commands │ └── __init__.py │ ├── _embedding │ ├── __init__.py │ ├── writers │ │ ├── __init__.py │ │ ├── writer_helpers.py │ │ └── torch_writer.py │ ├── embedding_format.py │ ├── embedding_transform.py │ └── embedding_predictor.py │ ├── _task_models │ ├── __init__.py │ ├── dinov2_eomt_semantic_segmentation │ │ ├── __init__.py │ │ ├── scale_block.py │ │ └── scheduler.py │ ├── dinov3_eomt_instance_segmentation │ │ ├── __init__.py │ │ ├── scale_block.py │ │ └── scheduler.py │ ├── dinov3_eomt_panoptic_segmentation │ │ ├── __init__.py │ │ ├── scale_block.py │ │ └── scheduler.py │ ├── dinov3_eomt_semantic_segmentation │ │ ├── __init__.py │ │ ├── scale_block.py │ │ └── scheduler.py │ ├── dinov3_ltdetr_object_detection │ │ └── dinov3_convnext_wrapper.py │ ├── dinov2_ltdetr_object_detection │ │ └── dinov2_vit_wrapper.py │ └── train_model_helpers.py │ ├── _transforms │ ├── __init__.py │ ├── predict_transform.py │ ├── task_transform.py │ ├── predict_semantic_segmentation_transform.py │ └── random_order.py │ ├── model_wrappers.py │ ├── errors.py │ ├── _task_checkpoint.py │ ├── _train_task_state.py │ ├── _float32_matmul_precision.py │ ├── _lightning_rank_zero.py │ └── _scaling.py ├── .gitattributes ├── docs ├── source │ ├── _templates │ │ ├── sidebar │ │ │ └── version.html │ │ ├── gtm │ │ │ ├── body.html │ │ │ └── head.html │ │ └── page.html │ ├── tutorials │ │ ├── yolo │ │ │ ├── results_VOC.png │ │ │ └── samples_VOC_train2012.png │ │ ├── depth_estimation │ │ │ ├── megadepth_samples.png │ │ │ └── diode_outdoor_samples.png │ │ ├── resnet │ │ │ ├── human-detection-dataset-fine-tuning.jpg │ │ │ └── human-detection-dataset-pretraining.jpg │ │ ├── index.md │ │ └── colab │ │ │ └── index.md │ ├── _static │ │ ├── images │ │ │ ├── object_detection │ │ │ │ └── street.jpg │ │ │ ├── instance_segmentation │ │ │ │ └── cats.jpg │ │ │ ├── panoptic_segmentation │ │ │ │ └── train.jpg │ │ │ └── tutorials │ │ │ │ └── embedding │ │ │ │ ├── umap_lightly_train.jpg │ │ │ │ ├── umap_lightly_train_colored.jpg │ │ │ │ └── umap_lightly_train_imagenet_colored.jpg │ │ └── custom.css │ ├── python_api │ │ ├── index.md │ │ └── lightly_train.md │ ├── data │ │ ├── index.md │ │ └── multi_channel.md │ ├── performance │ │ └── hardware_recommendations.md │ └── pretrain_distill │ │ └── models │ │ └── timm.md ├── format_code.py ├── Makefile └── build.py ├── dev_tools ├── dinov3_licenseheader.tmpl ├── licenseheader.tmpl ├── dinov2_licenseheader.tmpl ├── eomt_licenseheader.tmpl ├── deimv2_licenseheader.tmpl ├── rtdetr_licenseheader.tmpl └── pytorch_lightning_licenseheader.tmpl ├── tests ├── __init__.py ├── _data │ ├── __init__.py │ └── test_cache.py ├── _callbacks │ └── __init__.py ├── _commands │ ├── __init__.py │ └── test_train_task_helpers.py ├── _configs │ ├── __init__.py │ └── test_config.py ├── _embedding │ ├── __init__.py │ └── writers │ │ ├── __init__.py │ │ ├── test_writer_helpers.py │ │ └── test_torch_writer.py ├── _events │ └── __init__.py ├── _methods │ ├── __init__.py │ ├── detcon │ │ └── __init__.py │ ├── dino │ │ ├── __init__.py │ │ └── test_dino_transform.py │ ├── dinov2 │ │ ├── __init__.py │ │ └── test_dinov2_transform.py │ ├── simclr │ │ ├── __init__.py │ │ ├── test_simclr_transform.py │ │ └── test_simclr.py │ ├── densecl │ │ ├── __init__.py │ │ ├── test_densecl_transform.py │ │ └── test_densecl.py │ ├── distillation │ │ ├── __init__.py │ │ └── test_distillation_transform.py │ ├── distillationv2 │ │ ├── __init__.py │ │ └── test_distillationv2_transforms.py │ ├── test_method_args.py │ └── test_method_helpers.py ├── _models │ ├── __init__.py │ ├── custom │ │ └── __init__.py │ ├── dinov3 │ │ └── __init__.py │ ├── rfdetr │ │ └── __init__.py │ ├── timm │ │ └── __init__.py │ ├── dinov2_vit │ │ └── __init__.py │ ├── torchvision │ │ ├── __init__.py │ │ ├── test_resnet.py │ │ ├── test_torchvision_package.py │ │ ├── test_convnext.py │ │ └── test_shufflenet.py │ ├── super_gradients │ │ └── __init__.py │ ├── test_model_wrapper.py │ └── test_embedding_model.py ├── _optim │ └── __init__.py ├── _transforms │ ├── __init__.py │ ├── test_random_order.py │ └── test_normalize.py ├── test__plot.py ├── test__float32_matmul_precision.py ├── test__scaling.py └── test__checkpoint.py ├── CITATION.cff ├── docker ├── variables.mk ├── README.md └── Dockerfile-amd64-cuda ├── .github ├── ISSUE_TEMPLATE │ ├── docs_improvement.md │ ├── question.md │ ├── feature_request.md │ ├── setup_help.md │ └── bug_report.md ├── pull_request_template.md └── workflows │ ├── test_build.yml │ └── test_documentation.yml ├── .pre-commit-config.yaml ├── licences └── EOMT_LICENSE └── CONTRIBUTING.md /src/lightly_train/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # Don't add Jupyter Notebooks to GitHub's language statistics. 2 | *.ipynb linguist-vendored -------------------------------------------------------------------------------- /docs/source/_templates/sidebar/version.html: -------------------------------------------------------------------------------- 1 | {{ version }} ▼ 2 | -------------------------------------------------------------------------------- /docs/source/tutorials/yolo/results_VOC.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/tutorials/yolo/results_VOC.png -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vits14.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_small 3 | patch_size: 14 4 | drop_path_rate: 0.1 5 | -------------------------------------------------------------------------------- /docs/source/_static/images/object_detection/street.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/_static/images/object_detection/street.jpg -------------------------------------------------------------------------------- /docs/source/tutorials/yolo/samples_VOC_train2012.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/tutorials/yolo/samples_VOC_train2012.png -------------------------------------------------------------------------------- /docs/source/_static/images/instance_segmentation/cats.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/_static/images/instance_segmentation/cats.jpg -------------------------------------------------------------------------------- /docs/source/_static/images/panoptic_segmentation/train.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/_static/images/panoptic_segmentation/train.jpg -------------------------------------------------------------------------------- /docs/source/tutorials/depth_estimation/megadepth_samples.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/tutorials/depth_estimation/megadepth_samples.png -------------------------------------------------------------------------------- /docs/source/tutorials/depth_estimation/diode_outdoor_samples.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/tutorials/depth_estimation/diode_outdoor_samples.png -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vitb14.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_base 3 | patch_size: 14 4 | drop_path_rate: 0.2 5 | ffn_layer: swiglufused 6 | -------------------------------------------------------------------------------- /docs/source/_static/images/tutorials/embedding/umap_lightly_train.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/_static/images/tutorials/embedding/umap_lightly_train.jpg -------------------------------------------------------------------------------- /docs/source/tutorials/resnet/human-detection-dataset-fine-tuning.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/tutorials/resnet/human-detection-dataset-fine-tuning.jpg -------------------------------------------------------------------------------- /docs/source/tutorials/resnet/human-detection-dataset-pretraining.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/tutorials/resnet/human-detection-dataset-pretraining.jpg -------------------------------------------------------------------------------- /dev_tools/dinov3_licenseheader.tmpl: -------------------------------------------------------------------------------- 1 | Copyright (c) Meta Platforms, Inc. and affiliates. 2 | 3 | This software may be used and distributed in accordance with 4 | the terms of the DINOv3 License Agreement. -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/_vittest14.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: _vit_test 3 | patch_size: 14 4 | embed_dim: 8 5 | depth: 2 6 | num_heads: 2 7 | mlp_ratio: 1 -------------------------------------------------------------------------------- /docs/source/_static/images/tutorials/embedding/umap_lightly_train_colored.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/_static/images/tutorials/embedding/umap_lightly_train_colored.jpg -------------------------------------------------------------------------------- /dev_tools/licenseheader.tmpl: -------------------------------------------------------------------------------- 1 | Copyright (c) Lightly AG and affiliates. 2 | All rights reserved. 3 | 4 | This source code is licensed under the license found in the 5 | LICENSE file in the root directory of this source tree. 6 | -------------------------------------------------------------------------------- /docs/source/_static/images/tutorials/embedding/umap_lightly_train_imagenet_colored.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lightly-ai/lightly-train/HEAD/docs/source/_static/images/tutorials/embedding/umap_lightly_train_imagenet_colored.jpg -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /dev_tools/dinov2_licenseheader.tmpl: -------------------------------------------------------------------------------- 1 | Copyright (c) Meta Platforms, Inc. and affiliates. 2 | 3 | This source code is licensed under the Apache License, Version 2.0 4 | found in the LICENSE file in the root directory of this source tree. 5 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | -------------------------------------------------------------------------------- /tests/_data/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/hub/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | -------------------------------------------------------------------------------- /tests/_callbacks/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_commands/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_configs/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_embedding/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_events/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_optim/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/detcon/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/dino/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/dinov2/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/simclr/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/custom/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/dinov3/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/rfdetr/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/timm/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_transforms/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_configs/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_data/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_loggers/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vitb14_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_base 3 | patch_size: 14 4 | crops: 5 | global_crops_size: 518 # this is to set up the position embeddings properly 6 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vitl14_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_large 3 | patch_size: 14 4 | crops: 5 | global_crops_size: 518 # this is to set up the position embeddings properly 6 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vits14_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_small 3 | patch_size: 14 4 | crops: 5 | global_crops_size: 518 # this is to set up the position embeddings properly 6 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vits14_reg4.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_small 3 | patch_size: 14 4 | drop_path_rate: 0.1 5 | num_register_tokens: 4 6 | interpolate_antialias: true 7 | interpolate_offset: 0.0 8 | -------------------------------------------------------------------------------- /src/lightly_train/_optim/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_embedding/writers/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/densecl/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/dinov2_vit/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/torchvision/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /docs/source/python_api/index.md: -------------------------------------------------------------------------------- 1 | (python-api)= 2 | 3 | # Python API 4 | 5 | In the pages listed below, one can find the documentation for the Python modules of LightlyTrain. 6 | 7 | ```{toctree} 8 | :maxdepth: 1 9 | 10 | lightly_train 11 | ``` 12 | -------------------------------------------------------------------------------- /src/lightly_train/_callbacks/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_commands/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_embedding/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/dino/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/timm/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_transforms/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/distillation/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_methods/distillationv2/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /tests/_models/super_gradients/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | message: "If you use this software, please cite it as below." 3 | authors: 4 | - name: "Lightly Team" 5 | title: "LightlyTrain" 6 | type: software 7 | url: "https://github.com/lightly-ai/lightly-train" 8 | date-released: 2025 9 | -------------------------------------------------------------------------------- /src/lightly_train/_data/_serialize/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_embedding/writers/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/densecl/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/dinov2/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/simclr/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/custom/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/rfdetr/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/distillation/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/distillationv2/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/super_gradients/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/torchvision/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/ultralytics/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /dev_tools/eomt_licenseheader.tmpl: -------------------------------------------------------------------------------- 1 | --------------------------------------------------------------- 2 | © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 3 | Licensed under the MIT License. 4 | --------------------------------------------------------------- 5 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vitg14_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_giant2 3 | patch_size: 14 4 | ffn_layer: swiglufused 5 | crops: 6 | global_crops_size: 518 # this is to set up the position embeddings properly 7 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vitb14_reg4.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_base 3 | patch_size: 14 4 | drop_path_rate: 0.2 5 | ffn_layer: swiglufused 6 | num_register_tokens: 4 7 | interpolate_antialias: true 8 | interpolate_offset: 0.0 9 | -------------------------------------------------------------------------------- /docs/source/_templates/gtm/body.html: -------------------------------------------------------------------------------- 1 | 2 | 4 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov2_eomt_semantic_segmentation/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_instance_segmentation/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_panoptic_segmentation/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_semantic_segmentation/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/hub/utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | 7 | DINOV3_BASE_URL = "https://dl.fbaipublicfiles.com/dinov3" 8 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vitb14_reg4_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_base 3 | patch_size: 14 4 | num_register_tokens: 4 5 | interpolate_antialias: true 6 | interpolate_offset: 0.0 7 | crops: 8 | global_crops_size: 518 # this is to set up the position embeddings properly 9 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vitl14_reg4_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_large 3 | patch_size: 14 4 | num_register_tokens: 4 5 | interpolate_antialias: true 6 | interpolate_offset: 0.0 7 | crops: 8 | global_crops_size: 518 # this is to set up the position embeddings properly 9 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vits14_reg4_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_small 3 | patch_size: 14 4 | num_register_tokens: 4 5 | interpolate_antialias: true 6 | interpolate_offset: 0.0 7 | crops: 8 | global_crops_size: 518 # this is to set up the position embeddings properly 9 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/model_wrappers.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from lightly_train._models.rtdetr.rtdetr import RTDETRModelWrapper 9 | 10 | __all__ = ["RTDETRModelWrapper"] 11 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/eval/vitg14_reg4_pretrain.yaml: -------------------------------------------------------------------------------- 1 | student: 2 | arch: vit_giant2 3 | patch_size: 14 4 | ffn_layer: swiglufused 5 | num_register_tokens: 4 6 | interpolate_antialias: true 7 | interpolate_offset: 0.0 8 | crops: 9 | global_crops_size: 518 # this is to set up the position embeddings properly 10 | local_crops_size: 98 -------------------------------------------------------------------------------- /src/lightly_train/_optim/optimizer_type.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from enum import Enum 9 | 10 | 11 | class OptimizerType(Enum): 12 | ADAMW = "adamw" 13 | SGD = "sgd" 14 | LARS = "lars" 15 | -------------------------------------------------------------------------------- /docs/source/_templates/page.html: -------------------------------------------------------------------------------- 1 | {% extends "furo/page.html" %} 2 | 3 | {# Include GTM head code in the HTML head section #} 4 | {% block extrahead %} 5 | {% include "gtm/head.html" %} 6 | {{ super() }} 7 | {% endblock %} 8 | 9 | {# Include GTM body code right after the opening body tag (Furo provides the `body` block) #} 10 | {% block body %} 11 | {% include "gtm/body.html" %} 12 | {{ super() }} 13 | {% endblock %} -------------------------------------------------------------------------------- /src/lightly_train/_embedding/embedding_format.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from enum import Enum 9 | 10 | 11 | class EmbeddingFormat(Enum): 12 | CSV = "csv" 13 | LIGHTLY_CSV = "lightly_csv" 14 | TORCH = "torch" 15 | -------------------------------------------------------------------------------- /docs/source/_templates/gtm/head.html: -------------------------------------------------------------------------------- 1 | 2 | 7 | -------------------------------------------------------------------------------- /docker/variables.mk: -------------------------------------------------------------------------------- 1 | # This file must only contain variables and no targets. 2 | 3 | # Directory of this file 4 | THIS_FILE_DIR := $(dir $(lastword $(MAKEFILE_LIST))) 5 | 6 | TAG ?= $(shell git -C $(THIS_FILE_DIR) rev-parse HEAD) 7 | IMAGE ?= train 8 | DOCKER_BUILDKIT := 1 9 | 10 | VERSION := $(shell grep '__version__' $(THIS_FILE_DIR)/../src/lightly_train/__init__.py | sed -E 's/[^0-9.]//g') 11 | VERSION_X := $(shell echo $(VERSION) | cut -d. -f1) 12 | VERSION_XY := $(shell echo $(VERSION) | cut -d. -f1-2) 13 | -------------------------------------------------------------------------------- /src/lightly_train/_optim/trainable_modules.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from dataclasses import dataclass 9 | from typing import Iterable 10 | 11 | from torch.nn import Module 12 | 13 | 14 | @dataclass 15 | class TrainableModules: 16 | modules: Iterable[Module] 17 | modules_no_weight_decay: Iterable[Module] = () 18 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/docs_improvement.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: 📚 Documentation Improvement 3 | about: Report an issue or suggest improvement in documentation 4 | title: '[DOC] ' 5 | labels: documentation 6 | assignees: '' 7 | --- 8 | 9 | ### 📍 Location of the issue 10 | 11 | - File/Section: 12 | - Link (if applicable): 13 | 14 | ### ✏️ Suggested Improvement 15 | 16 | What could be improved, corrected, or added? 17 | 18 | ### 💬 Additional Notes 19 | 20 | Any context or references to clarify the documentation issue. 21 | -------------------------------------------------------------------------------- /docs/source/tutorials/index.md: -------------------------------------------------------------------------------- 1 | (tutorials)= 2 | 3 | # Tutorials 4 | 5 | In the following, some tutorials using LightlyTrain can be found. The tutorials are a good starting point to learn about LightlyTrain. 6 | 7 | ```{toctree} 8 | :maxdepth: 1 9 | 10 | resnet/index 11 | yolo/index 12 | depth_estimation/index 13 | embedding/index 14 | ``` 15 | 16 | We've also prepared some Google Colab notebooks for some packages we support. You can find them in the following page: 17 | 18 | ```{toctree} 19 | :maxdepth: 1 20 | colab/index 21 | ``` 22 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/question.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: ❓ Question 3 | about: Ask about how something works or project usage 4 | title: '[QUESTION] ' 5 | labels: question 6 | assignees: '' 7 | --- 8 | 9 | ### 🤔 What’s your question? 10 | 11 | Clearly state your question. 12 | 13 | ### 🧠 What did you try? 14 | 15 | Explain what you’ve already looked into or attempted. 16 | 17 | ### 📌 Related code or context (if any) 18 | 19 | Optional, but helpful for understanding your context. 20 | 21 | ```python 22 | # Related snippet if relevant 23 | ``` 24 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/kynan/nbstripout.git 3 | rev: 0.8.1 4 | hooks: 5 | - id: nbstripout 6 | description: "Strips outputs from Jupyter notebooks." 7 | - repo: https://github.com/pre-commit/pre-commit-hooks 8 | rev: v5.0.0 9 | hooks: 10 | - id: check-added-large-files 11 | description: "Avoid committing large files." 12 | args: ["--maxkb=2000", "--enforce-all"] 13 | - repo: https://github.com/fastai/nbdev 14 | rev: 2.4.2 15 | hooks: 16 | - id: nbdev_clean 17 | args: [--fname=examples] -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | 7 | # TODO(Lionel, 08/25): Remove the linter skip. 8 | # ruff: noqa 9 | 10 | from .dtype import as_torch_dtype 11 | from .utils import ( 12 | cat_keep_shapes, 13 | count_parameters, 14 | fix_random_seeds, 15 | get_conda_env, 16 | get_sha, 17 | named_apply, 18 | named_replace, 19 | uncat_with_shapes, 20 | ) 21 | -------------------------------------------------------------------------------- /dev_tools/deimv2_licenseheader.tmpl: -------------------------------------------------------------------------------- 1 | Licensed under the Apache License, Version 2.0 (the "License"); 2 | you may not use this file except in compliance with the License. 3 | You may obtain a copy of the License at 4 | 5 | http://www.apache.org/licenses/LICENSE-2.0 6 | 7 | Unless required by applicable law or agreed to in writing, software 8 | distributed under the License is distributed on an "AS IS" BASIS, 9 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 10 | See the License for the specific language governing permissions and 11 | limitations under the License. -------------------------------------------------------------------------------- /dev_tools/rtdetr_licenseheader.tmpl: -------------------------------------------------------------------------------- 1 | Licensed under the Apache License, Version 2.0 (the "License"); 2 | you may not use this file except in compliance with the License. 3 | You may obtain a copy of the License at 4 | 5 | http://www.apache.org/licenses/LICENSE-2.0 6 | 7 | Unless required by applicable law or agreed to in writing, software 8 | distributed under the License is distributed on an "AS IS" BASIS, 9 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 10 | See the License for the specific language governing permissions and 11 | limitations under the License. -------------------------------------------------------------------------------- /src/lightly_train/_models/torchvision/torchvision.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from torch.nn import Module 11 | 12 | from lightly_train._models.model_wrapper import ModelWrapper 13 | 14 | 15 | class TorchvisionModelWrapper(Module, ModelWrapper): 16 | _torchvision_models: list[type[Module]] 17 | # Regex pattern for matching model names. 18 | _torchvision_model_name_pattern: str 19 | -------------------------------------------------------------------------------- /dev_tools/pytorch_lightning_licenseheader.tmpl: -------------------------------------------------------------------------------- 1 | Copyright The Lightning AI team. 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: 🚀 Feature Request 3 | about: Suggest a new idea or enhancement 4 | title: '[FEAT] ' 5 | labels: enhancement 6 | assignees: '' 7 | --- 8 | 9 | ### 💡 Is your feature request related to a problem? 10 | 11 | A clear and concise description of the problem you're trying to solve. 12 | 13 | ### 🧰 Describe the solution you'd like 14 | 15 | What should the feature do? How would it work? 16 | 17 | ### 🛠 Alternatives you've considered 18 | 19 | Have you considered any alternative solutions? 20 | 21 | ### 📝 Additional context 22 | 23 | Add any other context, designs, or screenshots here. 24 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/setup_help.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: ⚙️ Environment Setup Help 3 | about: Need help with installing or running the project 4 | title: '[SETUP] ' 5 | labels: question, help wanted 6 | assignees: '' 7 | --- 8 | 9 | ### 🧪 What did you try? 10 | 11 | Steps you followed to set things up. 12 | 13 | ### 🧱 Error messages (if any) 14 | 15 | Paste any error logs or console output: 16 | 17 | ``` 18 | 19 | ``` 20 | 21 | ### 🤖 Environment Details 22 | 23 | - OS: 24 | - Python version:~ 25 | - Frameworks/Libraries (with versions): 26 | - How did you install the package: 27 | - Any unusual setups (Docker, cloud, WSL, etc.): 28 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vitg14.yaml: -------------------------------------------------------------------------------- 1 | dino: 2 | head_n_prototypes: 131072 3 | head_bottleneck_dim: 384 4 | ibot: 5 | separate_head: true 6 | head_n_prototypes: 131072 7 | train: 8 | batch_size_per_gpu: 12 9 | dataset_path: ImageNet22k 10 | centering: sinkhorn_knopp 11 | student: 12 | arch: vit_giant2 13 | patch_size: 14 14 | drop_path_rate: 0.4 15 | ffn_layer: swiglufused 16 | block_chunks: 4 17 | teacher: 18 | momentum_teacher: 0.994 19 | optim: 20 | epochs: 500 21 | weight_decay_end: 0.2 22 | base_lr: 2.0e-04 # learning rate for a batch size of 1024 23 | warmup_epochs: 80 24 | layerwise_decay: 1.0 25 | crops: 26 | local_crops_size: 98 27 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vitl14.yaml: -------------------------------------------------------------------------------- 1 | dino: 2 | head_n_prototypes: 131072 3 | head_bottleneck_dim: 384 4 | ibot: 5 | separate_head: true 6 | head_n_prototypes: 131072 7 | train: 8 | batch_size_per_gpu: 32 9 | dataset_path: ImageNet22k 10 | centering: sinkhorn_knopp 11 | student: 12 | arch: vit_large 13 | patch_size: 14 14 | drop_path_rate: 0.4 15 | ffn_layer: swiglufused 16 | block_chunks: 4 17 | teacher: 18 | momentum_teacher: 0.994 19 | optim: 20 | epochs: 500 21 | weight_decay_end: 0.2 22 | base_lr: 2.0e-04 # learning rate for a batch size of 1024 23 | warmup_epochs: 80 24 | layerwise_decay: 1.0 25 | crops: 26 | local_crops_size: 98 27 | -------------------------------------------------------------------------------- /src/lightly_train/errors.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | class LightlyTrainError(Exception): 9 | pass 10 | 11 | 12 | class UnknownModelError(LightlyTrainError): 13 | pass 14 | 15 | 16 | class ConfigError(LightlyTrainError): 17 | pass 18 | 19 | 20 | class ConfigUnknownKeyError(ConfigError): 21 | pass 22 | 23 | 24 | class ConfigValidationError(ConfigError): 25 | pass 26 | 27 | 28 | class ConfigMissingKeysError(ConfigError): 29 | pass 30 | 31 | 32 | class UnresolvedAutoError(LightlyTrainError): 33 | pass 34 | -------------------------------------------------------------------------------- /src/lightly_train/_models/super_gradients/super_gradients.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Protocol 11 | 12 | from torch.nn import Module 13 | 14 | from lightly_train._models.model_wrapper import ModelWrapper 15 | 16 | 17 | class SuperGradientsModelWrapper(ModelWrapper, Protocol): 18 | @classmethod 19 | def is_supported_model_cls(cls, model_cls: type[Module]) -> bool: ... 20 | 21 | @classmethod 22 | def supported_model_classes(cls) -> tuple[type[Module], ...]: ... 23 | -------------------------------------------------------------------------------- /.github/pull_request_template.md: -------------------------------------------------------------------------------- 1 | ## What has changed and why? 2 | 3 | (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.) 4 | 5 | ## How has it been tested? 6 | 7 | (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.) 8 | 9 | ## Did you update [CHANGELOG.md](../CHANGELOG.md)? 10 | 11 | - [ ] Yes 12 | - [ ] Not needed (internal change) 13 | 14 | ## Did you update the documentation? 15 | 16 | - [ ] Yes 17 | - [ ] Not needed (internal change without effects for user) 18 | -------------------------------------------------------------------------------- /docs/source/data/index.md: -------------------------------------------------------------------------------- 1 | (data-input)= 2 | 3 | # Data Input 4 | 5 | LightlyTrain supports standard RGB images and images for domain-specific use cases, such as medical imaging. It also loads images from specialized formats. 6 | 7 | The following image formats are supported: 8 | 9 | - jpg 10 | - jpeg 11 | - png 12 | - ppm 13 | - bmp 14 | - pgm 15 | - tif 16 | - tiff 17 | - webp 18 | - dcm [(DICOM)](#dicom-support) 19 | 20 | See the [Single- and Multi-Channel Images](multi_channel) page for details on using multi-channel images in LightlyTrain. 21 | 22 | See the dedicated page for more details about specific domain image support. 23 | 24 | ```{toctree} 25 | --- 26 | hidden: 27 | maxdepth: 1 28 | --- 29 | multi_channel 30 | dicom 31 | ``` 32 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vitg14_reg4.yaml: -------------------------------------------------------------------------------- 1 | dino: 2 | head_n_prototypes: 131072 3 | head_bottleneck_dim: 384 4 | ibot: 5 | separate_head: true 6 | head_n_prototypes: 131072 7 | train: 8 | batch_size_per_gpu: 12 9 | dataset_path: ImageNet22k 10 | centering: sinkhorn_knopp 11 | student: 12 | arch: vit_giant2 13 | patch_size: 14 14 | drop_path_rate: 0.4 15 | ffn_layer: swiglufused 16 | block_chunks: 4 17 | num_register_tokens: 4 18 | interpolate_antialias: true 19 | interpolate_offset: 0.0 20 | teacher: 21 | momentum_teacher: 0.994 22 | optim: 23 | epochs: 500 24 | weight_decay_end: 0.2 25 | base_lr: 2.0e-04 # learning rate for a batch size of 1024 26 | warmup_epochs: 80 27 | layerwise_decay: 1.0 28 | crops: 29 | local_crops_size: 98 30 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/configs/train/vitl14_reg4.yaml: -------------------------------------------------------------------------------- 1 | dino: 2 | head_n_prototypes: 131072 3 | head_bottleneck_dim: 384 4 | ibot: 5 | separate_head: true 6 | head_n_prototypes: 131072 7 | train: 8 | batch_size_per_gpu: 32 9 | dataset_path: ImageNet22k 10 | centering: sinkhorn_knopp 11 | student: 12 | arch: vit_large 13 | patch_size: 14 14 | drop_path_rate: 0.4 15 | ffn_layer: swiglufused 16 | block_chunks: 4 17 | num_register_tokens: 4 18 | interpolate_antialias: true 19 | interpolate_offset: 0.0 20 | teacher: 21 | momentum_teacher: 0.994 22 | optim: 23 | epochs: 500 24 | weight_decay_end: 0.2 25 | base_lr: 2.0e-04 # learning rate for a batch size of 1024 26 | warmup_epochs: 80 27 | layerwise_decay: 1.0 28 | crops: 29 | local_crops_size: 98 30 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: 🐞 Bug Report 3 | about: Report a problem or unexpected behavior 4 | title: '[BUG] ' 5 | labels: bug 6 | assignees: '' 7 | --- 8 | 9 | ### 🧠 Describe the Bug 10 | 11 | A clear and concise description of what the bug is. 12 | 13 | ### 🔁 Steps to Reproduce 14 | 15 | Please provide a minimal, self-contained code example that reproduces the issue: 16 | 17 | ```python 18 | # Minimal code example 19 | ``` 20 | 21 | ### 🤖 Environment Details 22 | 23 | - OS: 24 | - Python version: 25 | - Frameworks/Libraries (with versions): 26 | - How did you install the package: 27 | 28 | ### 📷 Screenshots (optional) 29 | 30 | If applicable, add screenshots to help explain your problem. 31 | 32 | ### 📌 Additional Context 33 | 34 | Add any other context about the problem here. 35 | -------------------------------------------------------------------------------- /docs/source/_static/custom.css: -------------------------------------------------------------------------------- 1 | /* Modern styling for both light and dark themes using Furo's CSS variables */ 2 | 3 | /* Sidebar version styling - works with both themes */ 4 | .sidebar-version { 5 | display: flex; 6 | color: var(--color-brand-primary); 7 | font-size: 1rem; 8 | text-align: left; 9 | padding: var(--sidebar-item-spacing-vertical) var(--sidebar-item-spacing-horizontal); 10 | } 11 | 12 | .sidebar-version:visited { 13 | color: var(--color-brand-primary); 14 | } 15 | 16 | /* Simple responsive banner classes */ 17 | .mobile-only { 18 | display: block; 19 | } 20 | 21 | .desktop-only { 22 | display: none; 23 | } 24 | 25 | @media (min-width: 768px) { 26 | .mobile-only { 27 | display: none; 28 | } 29 | 30 | .desktop-only { 31 | display: block; 32 | } 33 | } 34 | -------------------------------------------------------------------------------- /docs/source/tutorials/colab/index.md: -------------------------------------------------------------------------------- 1 | (colab)= 2 | 3 | # Google Colab Notebooks 4 | 5 | Here are some Google Colab notebooks that demonstrate how to use LightlyTrain with 6 | various packages. You can run these notebooks directly in Google Colab to see how 7 | LightlyTrain works in practice. 8 | 9 | - [Ultralytics](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/ultralytics_yolo.ipynb) 10 | - [YOLOv12](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/yolov12.ipynb) 11 | - [RF-DETR](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/rfdetr.ipynb) 12 | - [Torchvision Embedding Model](https://colab.research.google.com/github/lightly-ai/lightly-train/blob/main/examples/notebooks/torchvision_embedding_model.ipynb) 13 | -------------------------------------------------------------------------------- /src/lightly_train/_loggers/wandb.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Literal 11 | 12 | from pytorch_lightning.loggers import WandbLogger as LightningWandbLogger 13 | 14 | from lightly_train._configs.config import PydanticConfig 15 | 16 | 17 | class WandbLoggerArgs(PydanticConfig): 18 | name: str | None = None 19 | version: str | None = None 20 | offline: bool = False 21 | anonymous: bool | None = None 22 | project: str | None = None 23 | log_model: bool | Literal["all"] = False 24 | prefix: str = "" 25 | checkpoint_name: str | None = None 26 | 27 | 28 | # No customizations necessary. 29 | WandbLogger = LightningWandbLogger 30 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/layers/rms_norm.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | 7 | from __future__ import annotations 8 | 9 | import torch 10 | from torch import Tensor, nn 11 | 12 | 13 | class RMSNorm(nn.Module): 14 | def __init__(self, dim: int, eps: float = 1e-5): 15 | super().__init__() 16 | self.weight = nn.Parameter(torch.ones(dim)) 17 | self.eps = eps 18 | 19 | def reset_parameters(self) -> None: 20 | nn.init.constant_(self.weight, 1) 21 | 22 | def _norm(self, x: Tensor) -> Tensor: 23 | return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) 24 | 25 | def forward(self, x: Tensor) -> Tensor: 26 | output = self._norm(x.float()).type_as(x) 27 | return output * self.weight 28 | -------------------------------------------------------------------------------- /tests/_methods/test_method_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import pytest 9 | 10 | from lightly_train._methods import method_helpers 11 | from lightly_train._methods.method import Method 12 | 13 | _METHODS = method_helpers._list_methods() 14 | 15 | 16 | @pytest.mark.parametrize("method", _METHODS) 17 | def test_method_args__training_length_defaults(method: Method) -> None: 18 | method_args_cls = method_helpers.get_method_cls(method).method_args_cls() 19 | 20 | assert not ( 21 | method_args_cls.default_steps is None and method_args_cls.default_epochs is None 22 | ) 23 | assert not ( 24 | isinstance(method_args_cls.default_steps, int) 25 | and isinstance(method_args_cls.default_epochs, int) 26 | ) 27 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_ltdetr_object_detection/dinov3_convnext_wrapper.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from torch import Tensor 11 | from torch.nn import Module 12 | 13 | from lightly_train._models.dinov3.dinov3_src.models.convnext import ConvNeXt 14 | 15 | 16 | class DINOv3ConvNextWrapper(Module): 17 | def __init__(self, model: ConvNeXt) -> None: 18 | super().__init__() 19 | self.backbone = model 20 | 21 | def forward(self, x: Tensor) -> tuple[Tensor, ...]: 22 | feats = self.backbone.get_intermediate_layers(x, n=3, reshape=True) 23 | assert isinstance(feats, tuple) 24 | assert all(isinstance(f, Tensor) for f in feats) 25 | return feats 26 | -------------------------------------------------------------------------------- /src/lightly_train/_data/label_helpers.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from collections.abc import Iterable 11 | 12 | 13 | def get_class_id_to_internal_class_id_mapping( 14 | class_ids: Iterable[int], ignore_classes: set[int] | None 15 | ) -> dict[int, int]: 16 | """Returns mapping from class id to new class index. 17 | 18 | Skips ignored_classes. We use the class index internally as "internal class id" as 19 | some models require class ids to be in [0, num_classes - 1]. 20 | """ 21 | ignore_classes = ignore_classes or set() 22 | return { 23 | class_id: i 24 | for i, class_id in enumerate( 25 | class_id for class_id in class_ids if class_id not in ignore_classes 26 | ) 27 | } 28 | -------------------------------------------------------------------------------- /src/lightly_train/_transforms/predict_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pydantic import ConfigDict 11 | 12 | from lightly_train._configs.config import PydanticConfig 13 | from lightly_train.types import TransformInput, TransformOutput 14 | 15 | 16 | class PredictTransformArgs(PydanticConfig): 17 | model_config = ConfigDict(arbitrary_types_allowed=True) 18 | 19 | 20 | class PredictTransform: 21 | transform_args_cls: type[PredictTransformArgs] 22 | 23 | def __init__( 24 | self, 25 | transform_args: PredictTransformArgs, 26 | ) -> None: 27 | self.transform_args = transform_args 28 | 29 | def __call__(self, input: TransformInput) -> TransformOutput: 30 | raise NotImplementedError 31 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/layers/layer_scale.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the Apache License, Version 2.0 5 | # found in the LICENSE file in the root directory of this source tree. 6 | # 7 | 8 | # Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 9 | 10 | from typing import Union 11 | 12 | import torch 13 | from torch import Tensor, nn 14 | 15 | 16 | class LayerScale(nn.Module): 17 | def __init__( 18 | self, 19 | dim: int, 20 | init_values: Union[float, Tensor] = 1e-5, 21 | inplace: bool = False, 22 | ) -> None: 23 | super().__init__() 24 | self.inplace = inplace 25 | self.gamma = nn.Parameter(init_values * torch.ones(dim)) 26 | 27 | def forward(self, x: Tensor) -> Tensor: 28 | return x.mul_(self.gamma) if self.inplace else x * self.gamma 29 | -------------------------------------------------------------------------------- /src/lightly_train/_embedding/writers/writer_helpers.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from pathlib import Path 9 | 10 | from lightly_train._embedding.embedding_format import EmbeddingFormat 11 | from lightly_train._embedding.writers.csv_writer import CSVWriter 12 | from lightly_train._embedding.writers.embedding_writer import EmbeddingWriter 13 | from lightly_train._embedding.writers.torch_writer import TorchWriter 14 | 15 | 16 | def get_writer(format: EmbeddingFormat, filepath: Path) -> EmbeddingWriter: 17 | if CSVWriter.is_supported_format(format): 18 | return CSVWriter(filepath=filepath, format=format) 19 | elif TorchWriter.is_supported_format(format): 20 | return TorchWriter(filepath=filepath) 21 | else: 22 | raise ValueError(f"Unsupported embedding format: '{format}'") 23 | -------------------------------------------------------------------------------- /src/lightly_train/_loggers/logger_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pydantic import Field 11 | 12 | from lightly_train._configs.config import PydanticConfig 13 | from lightly_train._loggers.jsonl import JSONLLoggerArgs 14 | from lightly_train._loggers.mlflow import MLFlowLoggerArgs 15 | from lightly_train._loggers.tensorboard import TensorBoardLoggerArgs 16 | from lightly_train._loggers.wandb import WandbLoggerArgs 17 | 18 | 19 | class LoggerArgs(PydanticConfig): 20 | jsonl: JSONLLoggerArgs | None = Field(default_factory=JSONLLoggerArgs) 21 | mlflow: MLFlowLoggerArgs | None = None 22 | tensorboard: TensorBoardLoggerArgs | None = Field( 23 | default_factory=TensorBoardLoggerArgs 24 | ) 25 | wandb: WandbLoggerArgs | None = None 26 | -------------------------------------------------------------------------------- /src/lightly_train/_models/log_usage_example.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from os import linesep 11 | 12 | 13 | def format_log_msg_model_usage_example(log_message_code_block: list[str]) -> str: 14 | log_message_header = ( 15 | f"Example: How to use the exported model{linesep}{'-' * 88}{linesep}" 16 | ) 17 | 18 | log_message_footer = f"{'-' * 88}{linesep}" 19 | 20 | def format_code_lines(lines: list[str]) -> str: 21 | str_out = "" 22 | for line in lines: 23 | str_out += f"{line}{linesep}" 24 | return str_out 25 | 26 | log_message = ( 27 | log_message_header 28 | + format_code_lines(log_message_code_block) 29 | + log_message_footer 30 | ) 31 | 32 | return log_message 33 | -------------------------------------------------------------------------------- /src/lightly_train/_configs/omegaconf_utils.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Any 11 | 12 | from omegaconf import DictConfig, OmegaConf, SCMode 13 | 14 | 15 | def config_to_dict(config: DictConfig) -> dict[str, Any]: 16 | config_dict = OmegaConf.to_container( 17 | config, 18 | resolve=True, 19 | throw_on_missing=True, 20 | enum_to_str=False, 21 | structured_config_mode=SCMode.DICT, 22 | ) 23 | assert isinstance(config_dict, dict) 24 | # Type ignore required because OmegaConf.to_container() contains more possible 25 | # types than just dict[str, Any] but we know that it will always return string keys. 26 | result: dict[str, Any] = config_dict # type: ignore[assignment] 27 | return result 28 | -------------------------------------------------------------------------------- /tests/_methods/simclr/test_simclr_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import numpy as np 11 | 12 | from lightly_train._methods.simclr.simclr_transform import ( 13 | SimCLRTransform, 14 | ) 15 | from lightly_train.types import NDArrayImage, TransformInput 16 | 17 | 18 | def test_simclr_transform_shapes() -> None: 19 | img_np: NDArrayImage = np.random.uniform(0, 255, size=(1234, 1234, 3)).astype( 20 | np.uint8 21 | ) 22 | input: TransformInput = {"image": img_np} 23 | 24 | transform_args = SimCLRTransform.transform_args_cls()() 25 | transform = SimCLRTransform(transform_args) 26 | 27 | views = transform(input) 28 | assert len(views) == 2 29 | for view in views: 30 | assert view["image"].shape == (3, 224, 224) 31 | -------------------------------------------------------------------------------- /tests/_methods/densecl/test_densecl_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import numpy as np 11 | 12 | from lightly_train._methods.densecl.densecl_transform import ( 13 | DenseCLTransform, 14 | ) 15 | from lightly_train.types import NDArrayImage, TransformInput 16 | 17 | 18 | def test_densecl_transform_shapes() -> None: 19 | img_np: NDArrayImage = np.random.uniform(0, 255, size=(1234, 1234, 3)).astype( 20 | np.uint8 21 | ) 22 | input: TransformInput = {"image": img_np} 23 | 24 | transform_args = DenseCLTransform.transform_args_cls()() 25 | transform = DenseCLTransform(transform_args) 26 | 27 | views = transform(input) 28 | assert len(views) == 2 29 | for view in views: 30 | assert view["image"].shape == (3, 224, 224) 31 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/layers/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the Apache License, Version 2.0 5 | # found in the LICENSE file in the root directory of this source tree. 6 | # 7 | 8 | from lightly_train._models.dinov2_vit.dinov2_vit_src.layers.attention import ( 9 | MemEffAttention, 10 | ) 11 | from lightly_train._models.dinov2_vit.dinov2_vit_src.layers.block import ( 12 | NestedTensorBlock, 13 | ) 14 | from lightly_train._models.dinov2_vit.dinov2_vit_src.layers.mlp import Mlp 15 | from lightly_train._models.dinov2_vit.dinov2_vit_src.layers.patch_embed import ( 16 | PatchEmbed, 17 | ) 18 | from lightly_train._models.dinov2_vit.dinov2_vit_src.layers.swiglu_ffn import ( 19 | SwiGLUFFN, 20 | SwiGLUFFNFused, 21 | ) 22 | 23 | __all__ = [ 24 | "MemEffAttention", 25 | "NestedTensorBlock", 26 | "Mlp", 27 | "PatchEmbed", 28 | "SwiGLUFFN", 29 | "SwiGLUFFNFused", 30 | ] 31 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/layers/layer_scale.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | from __future__ import annotations 7 | 8 | from typing import Union 9 | 10 | import torch 11 | from torch import Tensor, nn 12 | 13 | 14 | class LayerScale(nn.Module): 15 | def __init__( 16 | self, 17 | dim: int, 18 | init_values: Union[float, Tensor] = 1e-5, 19 | inplace: bool = False, 20 | device=None, 21 | ) -> None: 22 | super().__init__() 23 | self.inplace = inplace 24 | self.gamma = nn.Parameter(torch.empty(dim, device=device)) 25 | self.init_values = init_values 26 | 27 | def reset_parameters(self): 28 | nn.init.constant_(self.gamma, self.init_values) 29 | 30 | def forward(self, x: Tensor) -> Tensor: 31 | return x.mul_(self.gamma) if self.inplace else x * self.gamma 32 | -------------------------------------------------------------------------------- /src/lightly_train/_optim/sgd_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from torch.optim.optimizer import Optimizer as TorchOptimizer 11 | from torch.optim.sgd import SGD 12 | 13 | from lightly_train._optim.optimizer_args import OptimizerArgs 14 | from lightly_train._optim.optimizer_type import OptimizerType 15 | from lightly_train.types import ParamsT 16 | 17 | 18 | class SGDArgs(OptimizerArgs): 19 | lr: float = 0.001 20 | momentum: float = 0.9 21 | weight_decay: float = 0.0001 22 | 23 | @staticmethod 24 | def type() -> OptimizerType: 25 | return OptimizerType.SGD 26 | 27 | def get_optimizer(self, params: ParamsT, lr_scale: float) -> TorchOptimizer: 28 | kwargs = self.model_dump() 29 | kwargs["lr"] *= lr_scale 30 | return SGD(params=params, **kwargs) 31 | -------------------------------------------------------------------------------- /src/lightly_train/_data/task_data_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pathlib import Path 11 | 12 | from lightly_train._configs.config import PydanticConfig 13 | from lightly_train._data.task_dataset import TaskDatasetArgs 14 | 15 | 16 | class TaskDataArgs(PydanticConfig): 17 | @property 18 | def included_classes(self) -> dict[int, str]: 19 | raise NotImplementedError() 20 | 21 | def train_imgs_path(self) -> Path: 22 | raise NotImplementedError() 23 | 24 | def val_imgs_path(self) -> Path: 25 | raise NotImplementedError() 26 | 27 | def get_train_args( 28 | self, 29 | ) -> TaskDatasetArgs: 30 | raise NotImplementedError() 31 | 32 | def get_val_args( 33 | self, 34 | ) -> TaskDatasetArgs: 35 | raise NotImplementedError() 36 | -------------------------------------------------------------------------------- /tests/_methods/dino/test_dino_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import numpy as np 11 | 12 | from lightly_train._methods.dino.dino_transform import DINOTransform 13 | from lightly_train.types import NDArrayImage, TransformInput 14 | 15 | 16 | def test_dino_transform_shapes() -> None: 17 | img_np: NDArrayImage = np.random.uniform(0, 255, size=(1234, 1234, 3)).astype( 18 | np.uint8 19 | ) 20 | input: TransformInput = {"image": img_np} 21 | 22 | transform_args = DINOTransform.transform_args_cls()() 23 | transform = DINOTransform(transform_args) 24 | 25 | views = transform(input) 26 | assert len(views) == 2 + 6 27 | for view in views[:2]: 28 | assert view["image"].shape == (3, 224, 224) 29 | for view in views[2:]: 30 | assert view["image"].shape == (3, 96, 96) 31 | -------------------------------------------------------------------------------- /src/lightly_train/_data/infinite_cycle_iterator.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from collections.abc import Iterable, Iterator 11 | from typing import Generic, TypeVar 12 | 13 | from typing_extensions import Self 14 | 15 | _T = TypeVar("_T") 16 | 17 | 18 | class InfiniteCycleIterator(Generic[_T]): 19 | def __init__(self, iterable: Iterable[_T]): 20 | self.iterable = iterable 21 | self.cycles = 0 22 | self._iter: Iterator[_T] | None = None 23 | 24 | def __iter__(self) -> Self: 25 | return self 26 | 27 | def __next__(self) -> _T: 28 | if self._iter is None: 29 | self._iter = iter(self.iterable) 30 | try: 31 | return next(self._iter) 32 | except StopIteration: 33 | self._iter = iter(self.iterable) 34 | self.cycles += 1 35 | return next(self._iter) 36 | -------------------------------------------------------------------------------- /tests/_commands/test_train_task_helpers.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import os 11 | from pathlib import Path 12 | 13 | import torch 14 | from lightning_fabric import Fabric 15 | 16 | from lightly_train._commands import train_task_helpers 17 | 18 | 19 | class DummyClass: 20 | pass 21 | 22 | 23 | def test__torch_weights_only_false(tmp_path: Path) -> None: 24 | fabric = Fabric(accelerator="cpu", devices=1) 25 | ckpt = {"dummy": DummyClass()} 26 | ckpt_path = tmp_path / "model.ckpt" 27 | fabric.save(ckpt_path, ckpt) # type: ignore 28 | assert os.environ.get("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD") is None 29 | with train_task_helpers._torch_weights_only_false(): 30 | assert os.environ.get("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD") == "1" 31 | torch.load(ckpt_path) 32 | fabric.load(ckpt_path) 33 | assert os.environ.get("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD") is None 34 | -------------------------------------------------------------------------------- /tests/_methods/distillation/test_distillation_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | import numpy as np 10 | from torch import Tensor 11 | 12 | from lightly_train._methods.distillation.distillation_transform import ( 13 | DistillationTransform, 14 | ) 15 | from lightly_train.types import TransformInput 16 | 17 | 18 | class TestDistillationTransform: 19 | def test_transform_shapes(self) -> None: 20 | img_np = np.random.uniform(0, 255, size=(1234, 1234, 3)) 21 | input: TransformInput = { 22 | "image": img_np.astype(np.uint8), 23 | } 24 | 25 | transform_args = DistillationTransform.transform_args_cls()() 26 | transform = DistillationTransform(transform_args) 27 | 28 | transformed = transform(input) 29 | assert len(transformed) == 1 30 | image = transformed[0]["image"] 31 | assert isinstance(image, Tensor) 32 | assert image.shape == (3, 224, 224) 33 | -------------------------------------------------------------------------------- /tests/test__plot.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import torch 9 | 10 | import lightly_train._plot as _plot 11 | from lightly_train.types import Batch 12 | 13 | 14 | def test_plot_example_augmentations() -> None: 15 | view_sizes = [(3, 128, 32), (3, 64, 64), (3, 64, 256)] 16 | n_images = 4 17 | expected_grid_height = 128 + 10 18 | expected_grid_width = 256 + 10 19 | expected_image_height = expected_grid_height * n_images + 30 20 | expected_image_width = expected_grid_width * len(view_sizes) + 100 21 | 22 | batch_per_view = [torch.rand(n_images, *size) for size in view_sizes] 23 | multi_view_batch: Batch = { 24 | "views": batch_per_view, 25 | "filename": [f"img_{i}" for i in range(n_images)], 26 | } 27 | 28 | pil_image = _plot.plot_example_augmentations( 29 | train_batch=multi_view_batch, max_examples=5 30 | ) 31 | assert pil_image.size == (expected_image_width, expected_image_height) 32 | -------------------------------------------------------------------------------- /tests/_methods/distillationv2/test_distillationv2_transforms.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | import numpy as np 10 | from torch import Tensor 11 | 12 | from lightly_train._methods.distillationv2.distillationv2_transform import ( 13 | DistillationV2Transform, 14 | ) 15 | from lightly_train.types import TransformInput 16 | 17 | 18 | class TestDistillationV2Transform: 19 | def test_transform_shapes(self) -> None: 20 | img_np = np.random.uniform(0, 255, size=(1234, 1234, 3)) 21 | input: TransformInput = { 22 | "image": img_np.astype(np.uint8), 23 | } 24 | 25 | transform_args = DistillationV2Transform.transform_args_cls()() 26 | transform = DistillationV2Transform(transform_args) 27 | 28 | transformed = transform(input) 29 | assert len(transformed) == 1 30 | image = transformed[0]["image"] 31 | assert isinstance(image, Tensor) 32 | assert image.shape == (3, 224, 224) 33 | -------------------------------------------------------------------------------- /src/lightly_train/_data/cache.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from pathlib import Path 9 | 10 | from lightly_train._env import Env 11 | 12 | 13 | def get_model_cache_dir() -> Path: 14 | """Returns the model cache directory for LightlyTrain, allowing override via env 15 | variable. 16 | """ 17 | # Get the cache directory from the environment variable if set. 18 | cache_dir = Env.LIGHTLY_TRAIN_MODEL_CACHE_DIR.value.expanduser().resolve() 19 | # Create the directory if it doesn't exist. 20 | cache_dir.mkdir(parents=True, exist_ok=True) 21 | return cache_dir 22 | 23 | 24 | def get_data_cache_dir() -> Path: 25 | """Get the data cache directory for LightlyTrain.""" 26 | # Get the cache directory from the environment variable if set. 27 | cache_dir = Env.LIGHTLY_TRAIN_DATA_CACHE_DIR.value.expanduser().resolve() 28 | # Create the directory if it doesn't exist. 29 | cache_dir.mkdir(parents=True, exist_ok=True) 30 | return cache_dir 31 | -------------------------------------------------------------------------------- /docs/format_code.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from pytest_examples import CodeExample, EvalExample, find_examples 3 | 4 | """ 5 | This is not a test, but provides the necessary commands to format and check 6 | python code in code blocks in the markdown files of the docs. 7 | It is called by `make format` and `make format-check`. 8 | 9 | Code blocks containing 'skip_ruff' are skipped. 10 | Additionally, as of `pytest-examples` v0.0.15 (2024-11-20), code blocks inside tabs 11 | are not detected (see https://github.com/pydantic/pytest-examples/issues/51). 12 | """ 13 | 14 | 15 | @pytest.mark.parametrize("example", find_examples("docs"), ids=str) 16 | def test_format_code_in_docs(example: CodeExample, eval_example: EvalExample) -> None: 17 | if "skip_ruff" in example.prefix: 18 | pytest.skip("Skip this example") 19 | eval_example.format_ruff(example) 20 | 21 | 22 | @pytest.mark.parametrize("example", find_examples("docs"), ids=str) 23 | def test_format_check_code_in_docs( 24 | example: CodeExample, eval_example: EvalExample 25 | ) -> None: 26 | if "skip_ruff" in example.prefix: 27 | pytest.skip("Skip this example") 28 | eval_example.lint_ruff(example) 29 | -------------------------------------------------------------------------------- /docs/source/performance/hardware_recommendations.md: -------------------------------------------------------------------------------- 1 | (hardware-recommendations)= 2 | 3 | # Hardware Recommendations 4 | 5 | We tested Lightly**Train** in various setups, the following training examples with hardware setups are provided as a performance reference: 6 | 7 | - **Distillation YOLOv8l on 2 × NVIDIA RTX 4090** 8 | 9 | - `method='distillation'` 10 | - `batch_size=512` 11 | - `model='ultralytics/yolov8l.yaml'` 12 | - `precision='bf16-mixed'` 13 | - Dataset Size: 1 million images 14 | - GPUs: 2 × NVIDIA RTX 4090 15 | - GPU Memory: ~16GB per GPU 16 | - Time per Epoch: ~19 minutes 17 | 18 | - **Distillation YOLO11x on 4 × NVIDIA H100** 19 | 20 | - `method='distillation'` 21 | - `batch_size=2048` 22 | - `model='ultralytics/yolov11x.yaml'` 23 | - `precision='32-true'` 24 | - Dataset Size: 1 million images 25 | - GPUs: 4 × NVIDIA H100 26 | - GPU Memory: ~80GB per GPU 27 | - Time per Epoch: ~6 minutes 28 | 29 | Those setups deliver efficient training while handling large-scale datasets with high throughput. While smaller setups can be used, training times may increase significantly and batch size might need to be reduced to fit GPU memory. 30 | -------------------------------------------------------------------------------- /licences/EOMT_LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2025 Mobile Perception Systems Lab at TU/e 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /src/lightly_train/_task_checkpoint.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import logging 9 | from typing import Literal 10 | 11 | from pydantic import model_validator 12 | from typing_extensions import Self 13 | 14 | from lightly_train._configs.config import PydanticConfig 15 | 16 | logger = logging.getLogger(__name__) 17 | 18 | 19 | class TaskSaveCheckpointArgs(PydanticConfig): 20 | save_every_num_steps: int = 1000 21 | save_last: bool = True 22 | save_best: bool = True 23 | watch_metric: str 24 | mode: Literal["min", "max"] 25 | 26 | @model_validator(mode="after") 27 | def _warn_if_no_checkpoints(self) -> Self: 28 | if not self.save_last and not self.save_best: 29 | logger.warning( 30 | "No checkpoints will be saved because both 'save_last' and 'save_best' " 31 | "are disabled. At least one of them should be enabled if checkpoint " 32 | "artifacts are required." 33 | ) 34 | return self 35 | -------------------------------------------------------------------------------- /src/lightly_train/_optim/optimizer_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | from torch.optim.optimizer import Optimizer 10 | 11 | from lightly_train._configs.config import PydanticConfig 12 | from lightly_train._optim.optimizer_type import OptimizerType 13 | from lightly_train.types import ParamsT 14 | 15 | 16 | class OptimizerArgs(PydanticConfig): 17 | """Base class for optimizer arguments.""" 18 | 19 | @staticmethod 20 | def type() -> OptimizerType: 21 | """Returns the optimizer type.""" 22 | raise NotImplementedError 23 | 24 | def get_optimizer(self, params: ParamsT, lr_scale: float) -> Optimizer: 25 | """Returns a new optimizer instance for the given parameters. 26 | 27 | Args: 28 | params: 29 | Parameters to optimize. 30 | lr_scale: 31 | Learning rate scale. Will be multiplied with the learning rate in the 32 | optimizer. 33 | """ 34 | raise NotImplementedError 35 | -------------------------------------------------------------------------------- /tests/_methods/dinov2/test_dinov2_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | # * Note: This file is almost identical to tests/_methods/dino/test_dino_transform.py 10 | from __future__ import annotations 11 | 12 | import numpy as np 13 | 14 | from lightly_train._methods.dinov2.dinov2_transform import ( 15 | DINOv2ViTTransform, 16 | ) 17 | from lightly_train.types import NDArrayImage, TransformInput 18 | 19 | 20 | def test_dinov2_transform_shapes() -> None: 21 | img_np: NDArrayImage = np.random.uniform(0, 255, size=(1234, 1234, 3)).astype( 22 | np.uint8 23 | ) 24 | input: TransformInput = {"image": img_np} 25 | 26 | transform_args = DINOv2ViTTransform.transform_args_cls()() 27 | transform = DINOv2ViTTransform(transform_args) 28 | 29 | views = transform(input) 30 | assert len(views) == 2 + 8 31 | for view in views[:2]: 32 | assert view["image"].shape == (3, 224, 224) 33 | for view in views[2:]: 34 | assert view["image"].shape == (3, 98, 98) 35 | -------------------------------------------------------------------------------- /src/lightly_train/_callbacks/learning_rate_monitor.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pytorch_lightning.callbacks import ( 11 | LearningRateMonitor as LightningLearningRateMonitor, 12 | ) 13 | 14 | 15 | class LearningRateMonitor(LightningLearningRateMonitor): 16 | def _get_optimizer_stats(self, *args, **kwargs) -> dict[str, float]: # type: ignore[no-untyped-def] 17 | # This fixes https://github.com/Lightning-AI/pytorch-lightning/issues/20250 18 | # The proper fix would be to add the float conversion in LightlySSL here: 19 | # https://github.com/lightly-ai/lightly/blob/ee30cd481d68862c80de4ef45920cfe1ab1f67b1/lightly/utils/scheduler.py#L67 20 | # But LightlyTrain has to be backwards compatible with older Lightly versions 21 | # so we add the fix here for now. 22 | stats = super()._get_optimizer_stats(*args, **kwargs) 23 | stats = {name: float(value) for name, value in stats.items()} 24 | return stats 25 | -------------------------------------------------------------------------------- /docker/README.md: -------------------------------------------------------------------------------- 1 | # LightlyTrain Docker Image 2 | 3 | Documentation on how to use the docker images: https://docs.lightly.ai/train/stable/docker.html 4 | 5 | ## Available Images 6 | 7 | List of currently available Docker base images: 8 | 9 | - `amd64-cuda` 10 | - More coming soon... 11 | 12 | TODO(Malte, 06/2024): Rethink and rework the setup of supporting different base images 13 | once we have multiple base images. Alternatives are e.g.: 14 | 15 | 1. Pass the base image type or directly the Dockerfile as argument to the makefile. 16 | 1. Put the Dockerfile, requirements and optionally makefile for each image type into 17 | a separate subdirectory. 18 | 1. Have docker multi-platform builds. 19 | 20 | ## Development 21 | 22 | ### Building Images 23 | 24 | Images are built by calling the corresponding [Makefile](./Makefile) command: 25 | 26 | - `make build-docker-IMAGE_TYPE` builds the image specified by the file `Dockerfile-IMAGE_TYPE` 27 | 28 | ### Testing Images 29 | 30 | Run tests with `make test`. The docker image must already be built. 31 | 32 | Attention! This requires that a Python environment is activated. It will also create 33 | some images locally (outside of the docker container) for testing. 34 | -------------------------------------------------------------------------------- /src/lightly_train/_optim/lars_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import cast 11 | 12 | from lightly.utils.lars import LARS 13 | from torch.optim.optimizer import Optimizer as TorchOptimizer 14 | 15 | from lightly_train._optim.optimizer_args import OptimizerArgs 16 | from lightly_train._optim.optimizer_type import OptimizerType 17 | from lightly_train.types import ParamsT 18 | 19 | 20 | class LARSArgs(OptimizerArgs): 21 | lr: float = 0.3 22 | momentum: float = 0 23 | dampening: float = 0 24 | weight_decay: float = 0 25 | nesterov: bool = False 26 | trust_coefficient: float = 0.001 27 | eps: float = 1e-8 28 | 29 | @staticmethod 30 | def type() -> OptimizerType: 31 | return OptimizerType.LARS 32 | 33 | def get_optimizer(self, params: ParamsT, lr_scale: float) -> TorchOptimizer: 34 | kwargs = self.model_dump() 35 | kwargs["lr"] *= lr_scale 36 | return cast(TorchOptimizer, LARS(params=params, **kwargs)) 37 | -------------------------------------------------------------------------------- /tests/_embedding/writers/test_writer_helpers.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pathlib import Path 11 | 12 | import pytest 13 | 14 | from lightly_train import EmbeddingFormat 15 | from lightly_train._embedding.writers import writer_helpers 16 | from lightly_train._embedding.writers.csv_writer import CSVWriter 17 | from lightly_train._embedding.writers.embedding_writer import EmbeddingWriter 18 | from lightly_train._embedding.writers.torch_writer import TorchWriter 19 | 20 | 21 | @pytest.mark.parametrize( 22 | "format, expected", 23 | [ 24 | (EmbeddingFormat.CSV, CSVWriter), 25 | (EmbeddingFormat.LIGHTLY_CSV, CSVWriter), 26 | (EmbeddingFormat.TORCH, TorchWriter), 27 | ], 28 | ) 29 | def test_get_writer( 30 | format: EmbeddingFormat, expected: type[EmbeddingWriter], tmp_path: Path 31 | ) -> None: 32 | filepath = tmp_path / "embeddings" 33 | writer = writer_helpers.get_writer(format=format, filepath=filepath) 34 | assert isinstance(writer, expected) 35 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/layers/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | 7 | # TODO(Lionel, 08/25): Remove the linting skip. 8 | # ruff: noqa 9 | 10 | from lightly_train._models.dinov3.dinov3_src.layers.attention import ( 11 | CausalSelfAttention, 12 | LinearKMaskedBias, 13 | SelfAttention, 14 | ) 15 | from lightly_train._models.dinov3.dinov3_src.layers.block import ( 16 | CausalSelfAttentionBlock, 17 | SelfAttentionBlock, 18 | ) 19 | from lightly_train._models.dinov3.dinov3_src.layers.ffn_layers import ( 20 | Mlp, 21 | SwiGLUFFN, 22 | ) 23 | from lightly_train._models.dinov3.dinov3_src.layers.fp8_linear import ( 24 | convert_linears_to_fp8, 25 | ) 26 | from lightly_train._models.dinov3.dinov3_src.layers.layer_scale import ( 27 | LayerScale, 28 | ) 29 | from lightly_train._models.dinov3.dinov3_src.layers.patch_embed import ( 30 | PatchEmbed, 31 | ) 32 | from lightly_train._models.dinov3.dinov3_src.layers.rms_norm import RMSNorm 33 | from lightly_train._models.dinov3.dinov3_src.layers.rope_position_encoding import ( 34 | RopePositionEmbedding, 35 | ) 36 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/method_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | from __future__ import annotations 10 | 11 | from typing import ClassVar, Literal 12 | 13 | from lightly_train._configs.config import PydanticConfig 14 | from lightly_train._models.model_wrapper import ModelWrapper 15 | from lightly_train._optim.optimizer_args import OptimizerArgs 16 | from lightly_train._scaling import ScalingInfo 17 | 18 | 19 | class MethodArgs(PydanticConfig): 20 | """Arguments for a method. 21 | 22 | This does not include optimizer or scheduler arguments. 23 | """ 24 | 25 | default_steps: ClassVar[int | None] = None 26 | default_epochs: ClassVar[int | None] = 100 27 | 28 | lr_scale_method: Literal["linear", "sqrt"] = "linear" 29 | reference_batch_size: int = 256 30 | 31 | def resolve_auto( 32 | self, 33 | scaling_info: ScalingInfo, 34 | optimizer_args: OptimizerArgs, 35 | wrapped_model: ModelWrapper, 36 | ) -> None: 37 | """Resolves all fields with the value 'auto' to their actual value.""" 38 | pass 39 | -------------------------------------------------------------------------------- /src/lightly_train/_train_task_state.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Any, TypedDict 11 | 12 | from torch.nn import Module 13 | from torch.optim.lr_scheduler import LRScheduler 14 | from torch.optim.optimizer import Optimizer 15 | from torch.utils.data import DataLoader 16 | 17 | from lightly_train.types import TaskBatch 18 | 19 | 20 | class TrainTaskState(TypedDict): 21 | train_model: Module 22 | optimizer: Optimizer 23 | scheduler: LRScheduler 24 | train_dataloader: DataLoader[TaskBatch] 25 | step: int 26 | # Model class path and initialization arguments for serialization. 27 | # Used to reconstruct the model after training. 28 | model_class_path: str 29 | model_init_args: dict[str, Any] 30 | 31 | 32 | class CheckpointDict(TypedDict): 33 | train_model_state_dict: dict[str, Any] 34 | # Model class path and initialization arguments for serialization. 35 | # Used to reconstruct the model after training. 36 | model_class_path: str 37 | model_init_args: dict[str, Any] 38 | -------------------------------------------------------------------------------- /src/lightly_train/_embedding/embedding_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import albumentations as A 11 | from albumentations.pytorch.transforms import ToTensorV2 12 | 13 | from lightly_train.types import ( 14 | TransformInput, 15 | TransformOutput, 16 | TransformOutputSingleView, 17 | ) 18 | 19 | 20 | class EmbeddingTransform: 21 | def __init__( 22 | self, 23 | image_size: int | tuple[int, int], 24 | mean: tuple[float, ...], 25 | std: tuple[float, ...], 26 | ): 27 | if isinstance(image_size, int): 28 | image_size = (image_size, image_size) 29 | self.transform = A.Compose( 30 | [ 31 | A.Resize(height=image_size[0], width=image_size[1]), 32 | A.Normalize(mean=mean, std=std, max_pixel_value=255.0), 33 | ToTensorV2(), 34 | ] 35 | ) 36 | 37 | def __call__(self, input: TransformInput) -> TransformOutput: 38 | transformed: TransformOutputSingleView = self.transform(**input) 39 | return [transformed] 40 | -------------------------------------------------------------------------------- /src/lightly_train/_optim/adamw_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pydantic import Field 11 | from torch.optim.adamw import AdamW 12 | from torch.optim.optimizer import Optimizer as TorchOptimizer 13 | 14 | from lightly_train._optim.optimizer_args import OptimizerArgs 15 | from lightly_train._optim.optimizer_type import OptimizerType 16 | from lightly_train.types import ParamsT 17 | 18 | 19 | class AdamWArgs(OptimizerArgs): 20 | lr: float = 0.001 21 | # Strict is set to False because OmegaConf does not support parsing tuples from the 22 | # CLI. Setting strict to False allows Pydantic to convert lists to tuples. 23 | betas: tuple[float, float] = Field(default=(0.9, 0.999), strict=False) 24 | eps: float = 1e-8 25 | weight_decay: float = 0.01 26 | 27 | @staticmethod 28 | def type() -> OptimizerType: 29 | return OptimizerType.ADAMW 30 | 31 | def get_optimizer(self, params: ParamsT, lr_scale: float) -> TorchOptimizer: 32 | kwargs = self.model_dump() 33 | kwargs["lr"] *= lr_scale 34 | return AdamW(params=params, **kwargs) 35 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/dinov2/scheduler.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | # References: 10 | # - https://github.com/lightly-ai/lightly/blob/master/lightly/utils/scheduler.py 11 | 12 | 13 | def linear_warmup_schedule( 14 | step: int, 15 | warmup_steps: int, 16 | start_value: float, 17 | end_value: float, 18 | ) -> float: # TODO: import from LightlySSL after new release 19 | if warmup_steps < 0: 20 | raise ValueError(f"Warmup steps {warmup_steps} can't be negative.") 21 | if step < 0: 22 | raise ValueError(f"Current step number {step} can't be negative.") 23 | if start_value < 0: 24 | raise ValueError(f"Start value {start_value} can't be negative.") 25 | if end_value <= 0: 26 | raise ValueError(f"End value {end_value} can't be non-positive.") 27 | if start_value > end_value: 28 | raise ValueError( 29 | f"Start value {start_value} must be less than or equal to end value {end_value}." 30 | ) 31 | if step < warmup_steps: 32 | return start_value + step / warmup_steps * (end_value - start_value) 33 | else: 34 | return end_value 35 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/utils/dtype.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | 7 | from __future__ import annotations 8 | 9 | from typing import Dict, Union 10 | 11 | import numpy as np 12 | import torch 13 | 14 | TypeSpec = Union[str, np.dtype, torch.dtype] 15 | 16 | 17 | _NUMPY_TO_TORCH_DTYPE: Dict[np.dtype, torch.dtype] = { 18 | np.dtype("bool"): torch.bool, 19 | np.dtype("uint8"): torch.uint8, 20 | np.dtype("int8"): torch.int8, 21 | np.dtype("int16"): torch.int16, 22 | np.dtype("int32"): torch.int32, 23 | np.dtype("int64"): torch.int64, 24 | np.dtype("float16"): torch.float16, 25 | np.dtype("float32"): torch.float32, 26 | np.dtype("float64"): torch.float64, 27 | np.dtype("complex64"): torch.complex64, 28 | np.dtype("complex128"): torch.complex128, 29 | } 30 | 31 | 32 | def as_torch_dtype(dtype: TypeSpec) -> torch.dtype: 33 | if isinstance(dtype, torch.dtype): 34 | return dtype 35 | if isinstance(dtype, str): 36 | dtype = np.dtype(dtype) 37 | assert isinstance(dtype, np.dtype), ( 38 | f"Expected an instance of nunpy dtype, got {type(dtype)}" 39 | ) 40 | return _NUMPY_TO_TORCH_DTYPE[dtype] 41 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov2_ltdetr_object_detection/dinov2_vit_wrapper.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from typing import Sequence, Tuple 9 | 10 | from torch import Tensor 11 | from torch.nn import Module 12 | 13 | from lightly_train._models.dinov2_vit.dinov2_vit_src.models.vision_transformer import ( 14 | DinoVisionTransformer, 15 | ) 16 | 17 | 18 | class DINOv2ViTWrapper(Module): 19 | # TODO: Lionel(09/25) Try the DEIMv2 wrapper: https://github.com/Intellindust-AI-Lab/DEIMv2/blob/main/engine/backbone/dinov3_adapter.py#L72 20 | def __init__( 21 | self, model: DinoVisionTransformer, keep_indices: Sequence[int] = (5, 8, 11) 22 | ): 23 | super().__init__() 24 | self.keep_indices = list(keep_indices) 25 | self.backbone = model 26 | 27 | def forward(self, x: Tensor) -> Tuple[Tensor, ...]: 28 | # TODO: Lionel(09/25) Infer minimum n from keep_indices. 29 | feats = self.backbone.get_intermediate_layers(x, n=12, reshape=True) 30 | feats_: list[Tensor] = [feats[i] for i in self.keep_indices] # type: ignore[misc] 31 | assert all(isinstance(f, Tensor) for f in feats_) 32 | return tuple(feats_) 33 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/train_model_helpers.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Any 11 | 12 | import torch 13 | from torch.nn import Module 14 | 15 | 16 | def criterion_empty_weight_reinit_hook( 17 | module: Module, 18 | state_dict: dict[str, Any], 19 | prefix: str, 20 | *args: Any, 21 | **kwargs: Any, 22 | ) -> None: 23 | criterion_empty_weight_key = f"{prefix}criterion.empty_weight" 24 | criterion_empty_weight = state_dict.get(criterion_empty_weight_key) 25 | if criterion_empty_weight is None: 26 | return 27 | 28 | criterion_module = getattr(module, "criterion", None) 29 | if criterion_module is None: 30 | return 31 | 32 | model_args_module = getattr(module, "model_args", None) 33 | if model_args_module is None: 34 | return 35 | 36 | # Re-initialize the empty weight buffer to match the current 37 | criterion_empty_weight_reinit = torch.ones_like(criterion_module.empty_weight) 38 | criterion_empty_weight_reinit[-1] = model_args_module.loss_no_object_coefficient 39 | 40 | state_dict[criterion_empty_weight_key] = criterion_empty_weight_reinit 41 | -------------------------------------------------------------------------------- /tests/test__float32_matmul_precision.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Literal 11 | 12 | import pytest 13 | import torch 14 | 15 | from lightly_train import _float32_matmul_precision 16 | 17 | 18 | @pytest.mark.parametrize( 19 | "precision, expected", 20 | [ 21 | ("auto", {"highest", "high", "medium"}), 22 | ("highest", {"highest"}), 23 | ("high", {"high"}), 24 | ("medium", {"medium"}), 25 | ], 26 | ) 27 | def test_get_float32_matmul_precision__auto( 28 | precision: Literal["auto", "highest", "high", "medium"], expected: set[str] 29 | ) -> None: 30 | assert _float32_matmul_precision.get_float32_matmul_precision(precision) in expected 31 | 32 | 33 | @pytest.mark.parametrize("precision", ["highest", "high", "medium"]) 34 | def test_float32_matmul_precision( 35 | precision: Literal["highest", "high", "medium"], 36 | ) -> None: 37 | default = torch.get_float32_matmul_precision() 38 | with _float32_matmul_precision.float32_matmul_precision(precision): 39 | assert torch.get_float32_matmul_precision() == precision 40 | assert torch.get_float32_matmul_precision() == default 41 | -------------------------------------------------------------------------------- /src/lightly_train/_callbacks/callback_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pydantic import Field 11 | 12 | from lightly_train._callbacks.checkpoint import ModelCheckpointArgs 13 | from lightly_train._callbacks.export import ModelExportArgs 14 | from lightly_train._configs.config import PydanticConfig 15 | 16 | 17 | class LearningRateMonitorArgs(PydanticConfig): 18 | pass 19 | 20 | 21 | class DeviceStatsMonitorArgs(PydanticConfig): 22 | pass 23 | 24 | 25 | class EarlyStoppingArgs(PydanticConfig): 26 | monitor: str = "train_loss" 27 | patience: int = int(1e12) 28 | check_finite: bool = True 29 | 30 | 31 | class CallbackArgs(PydanticConfig): 32 | learning_rate_monitor: LearningRateMonitorArgs | None = Field( 33 | default_factory=LearningRateMonitorArgs 34 | ) 35 | device_stats_monitor: DeviceStatsMonitorArgs | None = None 36 | early_stopping: EarlyStoppingArgs | None = Field(default_factory=EarlyStoppingArgs) 37 | model_export: ModelExportArgs | None = Field(default_factory=ModelExportArgs) 38 | model_checkpoint: ModelCheckpointArgs | None = Field( 39 | default_factory=ModelCheckpointArgs 40 | ) 41 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/layers/drop_path.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the Apache License, Version 2.0 5 | # found in the LICENSE file in the root directory of this source tree. 6 | # 7 | 8 | # References: 9 | # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py 10 | # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py 11 | 12 | 13 | from torch import nn 14 | 15 | 16 | def drop_path(x, drop_prob: float = 0.0, training: bool = False): 17 | if drop_prob == 0.0 or not training: 18 | return x 19 | keep_prob = 1 - drop_prob 20 | shape = (x.shape[0],) + (1,) * ( 21 | x.ndim - 1 22 | ) # work with diff dim tensors, not just 2D ConvNets 23 | random_tensor = x.new_empty(shape).bernoulli_(keep_prob) 24 | if keep_prob > 0.0: 25 | random_tensor.div_(keep_prob) 26 | output = x * random_tensor 27 | return output 28 | 29 | 30 | class DropPath(nn.Module): 31 | """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" 32 | 33 | def __init__(self, drop_prob=None): 34 | super(DropPath, self).__init__() 35 | self.drop_prob = drop_prob 36 | 37 | def forward(self, x): 38 | return drop_path(x, self.drop_prob, self.training) 39 | -------------------------------------------------------------------------------- /src/lightly_train/_models/rtdetr/rtdetr.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from torch import Tensor 11 | from torch.nn import AdaptiveAvgPool2d, Module 12 | 13 | from lightly_train._models.model_wrapper import ( 14 | ForwardFeaturesOutput, 15 | ForwardPoolOutput, 16 | ModelWrapper, 17 | ) 18 | 19 | 20 | class RTDETRModelWrapper(Module, ModelWrapper): 21 | def __init__(self, model: Module): 22 | super().__init__() 23 | self._model = [model] 24 | self._backbone = self._model[0].backbone 25 | self._pool = AdaptiveAvgPool2d((1, 1)) 26 | 27 | def get_model(self) -> Module: 28 | return self._model[0] 29 | 30 | def forward_features(self, x: Tensor) -> ForwardFeaturesOutput: 31 | features = self._backbone(x)[-1] # type: ignore[operator] 32 | return {"features": features} 33 | 34 | def forward_pool(self, x: ForwardFeaturesOutput) -> ForwardPoolOutput: 35 | return {"pooled_features": self._pool(x["features"])} 36 | 37 | def feature_dim(self) -> int: 38 | feat_dim = self._backbone.out_channels[-1] # type: ignore 39 | assert isinstance(feat_dim, int) 40 | return feat_dim 41 | -------------------------------------------------------------------------------- /src/lightly_train/_data/download.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import logging 9 | import urllib.request 10 | from pathlib import Path 11 | from urllib.error import HTTPError, URLError 12 | 13 | from tqdm import tqdm 14 | 15 | logger = logging.getLogger(__name__) 16 | 17 | 18 | def download_from_url(url: str, destination: Path, timeout: float = 10.0) -> None: 19 | """ 20 | Downloads a file from `url` to `destination` with timeout, progress bar, and error handling. 21 | """ 22 | try: 23 | with urllib.request.urlopen(url, timeout=timeout) as response: 24 | total = response.length 25 | with tqdm( 26 | total=total, unit="B", unit_scale=True, desc="Downloading", ncols=80 27 | ) as t: 28 | with open(destination, "wb") as out_file: 29 | while True: 30 | chunk = response.read(8192) 31 | if not chunk: 32 | break 33 | out_file.write(chunk) 34 | t.update(len(chunk)) 35 | except (URLError, HTTPError, TimeoutError, Exception) as e: 36 | raise RuntimeError(f"Download from '{url}' failed: '{e}'") 37 | -------------------------------------------------------------------------------- /src/lightly_train/_loggers/tensorboard.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import os 11 | 12 | from pytorch_lightning.loggers import TensorBoardLogger as LightningTensorBoardLogger 13 | from pytorch_lightning.utilities import rank_zero_only 14 | from typing_extensions import override 15 | 16 | from lightly_train._configs.config import PydanticConfig 17 | 18 | 19 | class TensorBoardLoggerArgs(PydanticConfig): 20 | name: str = "" 21 | version: str = "" 22 | log_graph: bool = False 23 | default_hp_metric: bool = True 24 | prefix: str = "" 25 | sub_dir: str | None = None 26 | 27 | 28 | class TensorBoardLogger(LightningTensorBoardLogger): 29 | @override 30 | @rank_zero_only # type: ignore[misc] 31 | def save(self) -> None: 32 | super().save() 33 | # Delete hparams file as the parent class creates it and there is no easy way 34 | # to disable it. We don't want hparams file because we'll save all parameters to 35 | # a custom file. 36 | dir_path = self.log_dir 37 | hparams_file = os.path.join(dir_path, self.NAME_HPARAMS_FILE) 38 | if self._fs.isfile(hparams_file): 39 | self._fs.rm(hparams_file) 40 | -------------------------------------------------------------------------------- /src/lightly_train/_float32_matmul_precision.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import contextlib 9 | from typing import Generator, Literal 10 | 11 | import torch 12 | 13 | 14 | def get_float32_matmul_precision( 15 | float32_matmul_precision: Literal["auto", "highest", "high", "medium"], 16 | ) -> Literal["highest", "high", "medium"]: 17 | """Get the float32 matmul precision setting.""" 18 | if float32_matmul_precision == "auto": 19 | # Return torch default precision or the value set by the user if they set it 20 | # with torch.set_float32_matmul_precision before. 21 | return torch.get_float32_matmul_precision() # type: ignore[return-value] 22 | else: 23 | return float32_matmul_precision 24 | 25 | 26 | @contextlib.contextmanager 27 | def float32_matmul_precision( 28 | float32_matmul_precision: Literal["highest", "high", "medium"], 29 | ) -> Generator[None, None, None]: 30 | """Context manager to temporarily set the float32 matmul precision.""" 31 | current_precision = torch.get_float32_matmul_precision() 32 | try: 33 | torch.set_float32_matmul_precision(float32_matmul_precision) 34 | yield 35 | finally: 36 | torch.set_float32_matmul_precision(current_precision) 37 | -------------------------------------------------------------------------------- /src/lightly_train/_embedding/embedding_predictor.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pytorch_lightning import LightningModule 11 | from torch import Tensor 12 | from torch.nn import Flatten 13 | 14 | from lightly_train._models.embedding_model import EmbeddingModel 15 | from lightly_train.types import Batch 16 | 17 | 18 | class EmbeddingPredictor(LightningModule): 19 | """PyTorch Lightning module for "predicting" embeddings. 20 | 21 | This module uses the `predict_step` to extract embeddings from the given 22 | embedding model. 23 | 24 | Args: 25 | embedding_model: The embedding model. 26 | """ 27 | 28 | def __init__(self, embedding_model: EmbeddingModel): 29 | super().__init__() 30 | self.embedding_model = embedding_model 31 | self.flatten = Flatten(start_dim=1) 32 | 33 | def forward(self, x: Tensor) -> Tensor: 34 | x = self.embedding_model(x) 35 | x = self.flatten(x) 36 | return x 37 | 38 | def predict_step(self, batch: Batch, batch_idx: int) -> tuple[Tensor, list[str]]: 39 | x = batch["views"][0] 40 | filenames = batch["filename"] 41 | embeddings = self(x) 42 | return embeddings, filenames 43 | -------------------------------------------------------------------------------- /src/lightly_train/_lightning_rank_zero.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright The Lightning AI team. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # 16 | 17 | from __future__ import annotations 18 | 19 | import os 20 | 21 | 22 | def get_global_rank() -> int | None: 23 | """Get the global rank of the current process. 24 | 25 | Implementation copied from: https://github.com/Lightning-AI/pytorch-lightning/blob/06a8d5bf33faf0a4f9a24207ae77b439354350af/src/lightning/fabric/utilities/rank_zero.py#L39-L49 26 | """ 27 | # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing, 28 | # therefore LOCAL_RANK needs to be checked first 29 | rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK") 30 | for key in rank_keys: 31 | rank = os.environ.get(key) 32 | if rank is not None: 33 | return int(rank) 34 | # None to differentiate whether an environment variable was set at all 35 | return None 36 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/layers/mlp.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the Apache License, Version 2.0 5 | # found in the LICENSE file in the root directory of this source tree. 6 | # 7 | 8 | # References: 9 | # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py 10 | # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py 11 | 12 | 13 | from typing import Callable, Optional 14 | 15 | from torch import Tensor, nn 16 | 17 | 18 | class Mlp(nn.Module): 19 | def __init__( 20 | self, 21 | in_features: int, 22 | hidden_features: Optional[int] = None, 23 | out_features: Optional[int] = None, 24 | act_layer: Callable[..., nn.Module] = nn.GELU, 25 | drop: float = 0.0, 26 | bias: bool = True, 27 | ) -> None: 28 | super().__init__() 29 | out_features = out_features or in_features 30 | hidden_features = hidden_features or in_features 31 | self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) 32 | self.act = act_layer() 33 | self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) 34 | self.drop = nn.Dropout(drop) 35 | 36 | def forward(self, x: Tensor) -> Tensor: 37 | x = self.fc1(x) 38 | x = self.act(x) 39 | x = self.drop(x) 40 | x = self.fc2(x) 41 | x = self.drop(x) 42 | return x 43 | -------------------------------------------------------------------------------- /tests/_transforms/test_random_order.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import numpy as np 9 | import pytest 10 | from albumentations import ColorJitter, Resize 11 | 12 | from lightly_train._transforms.random_order import RandomOrder 13 | 14 | 15 | class TestRandomOrder: 16 | def test__shapes(self) -> None: 17 | tr = RandomOrder( 18 | transforms=[ 19 | Resize(32, 32, p=1.0), 20 | ColorJitter(brightness=0.5, contrast=0.5), 21 | ], 22 | n=2, 23 | p=1.0, 24 | ) 25 | img = np.random.randint(0, 256, (8, 8, 3), dtype=np.uint8) 26 | out = tr(image=img) 27 | assert out["image"].shape == (32, 32, 3) 28 | 29 | @pytest.mark.parametrize("n", [0, 1, 2]) 30 | def test__get_idx(self, n: int) -> None: 31 | tr = RandomOrder( 32 | transforms=[ 33 | Resize(32, 32, p=1.0), 34 | ColorJitter(brightness=0.5, contrast=0.5), 35 | ], 36 | n=n, 37 | p=1.0, 38 | ) 39 | if n == 0: 40 | assert set(tr._get_idx().tolist()) == set() 41 | elif n == 1: 42 | assert set(tr._get_idx().tolist()).issubset({0, 1}) 43 | elif n == 2: 44 | assert set(tr._get_idx().tolist()) == {0, 1} 45 | -------------------------------------------------------------------------------- /src/lightly_train/_loggers/task_logger_args.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Literal 11 | 12 | from pydantic import Field 13 | 14 | from lightly_train._configs.config import PydanticConfig 15 | from lightly_train._loggers.mlflow import MLFlowLoggerArgs 16 | from lightly_train._loggers.tensorboard import TensorBoardLoggerArgs 17 | from lightly_train._loggers.wandb import WandbLoggerArgs 18 | 19 | 20 | class TaskLoggerArgs(PydanticConfig): 21 | log_every_num_steps: int | Literal["auto"] = "auto" 22 | val_every_num_steps: int | Literal["auto"] = "auto" 23 | val_log_every_num_steps: int | Literal["auto"] = "auto" 24 | 25 | mlflow: MLFlowLoggerArgs | None = None 26 | tensorboard: TensorBoardLoggerArgs | None = Field( 27 | default_factory=TensorBoardLoggerArgs 28 | ) 29 | wandb: WandbLoggerArgs | None = None 30 | 31 | def resolve_auto(self, steps: int, val_steps: int) -> None: 32 | if self.log_every_num_steps == "auto": 33 | self.log_every_num_steps = min(100, max(1, steps // 10)) 34 | if self.val_every_num_steps == "auto": 35 | self.val_every_num_steps = min(1000, max(1, steps)) 36 | if self.val_log_every_num_steps == "auto": 37 | self.val_log_every_num_steps = min(20, max(1, val_steps)) 38 | -------------------------------------------------------------------------------- /src/lightly_train/_models/torchvision/resnet.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from torch import Tensor 9 | from torchvision.models import ResNet 10 | from torchvision.models._utils import IntermediateLayerGetter 11 | 12 | from lightly_train._models.model_wrapper import ( 13 | ForwardFeaturesOutput, 14 | ForwardPoolOutput, 15 | ) 16 | from lightly_train._models.torchvision.torchvision import TorchvisionModelWrapper 17 | 18 | 19 | class ResNetModelWrapper(TorchvisionModelWrapper): 20 | _torchvision_models = [ResNet] 21 | _torchvision_model_name_pattern = r"resnet.*" 22 | 23 | def __init__(self, model: ResNet) -> None: 24 | super().__init__() 25 | self._model = [model] 26 | self._features = IntermediateLayerGetter( 27 | model=model, return_layers={"layer4": "out"} 28 | ) 29 | self._pool = model.avgpool 30 | self._feature_dim: int = model.fc.in_features 31 | 32 | def feature_dim(self) -> int: 33 | return self._feature_dim 34 | 35 | def forward_features(self, x: Tensor) -> ForwardFeaturesOutput: 36 | return {"features": self._features(x)["out"]} 37 | 38 | def forward_pool(self, x: ForwardFeaturesOutput) -> ForwardPoolOutput: 39 | return {"pooled_features": self._pool(x["features"])} 40 | 41 | def get_model(self) -> ResNet: 42 | return self._model[0] 43 | -------------------------------------------------------------------------------- /src/lightly_train/_transforms/task_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Any, TypedDict 11 | 12 | from pydantic import ConfigDict 13 | 14 | from lightly_train._configs.config import PydanticConfig 15 | 16 | 17 | class TaskTransformInput(TypedDict): 18 | pass 19 | 20 | 21 | class TaskTransformOutput(TypedDict): 22 | pass 23 | 24 | 25 | class TaskTransformArgs(PydanticConfig): 26 | def resolve_auto(self, model_init_args: dict[str, Any]) -> None: 27 | """Resolve any arguments set to "auto".""" 28 | pass 29 | 30 | def resolve_incompatible(self) -> None: 31 | """Resolve any incompatible arguments.""" 32 | pass 33 | 34 | model_config = ConfigDict(arbitrary_types_allowed=True) 35 | 36 | 37 | class TaskTransform: 38 | transform_args_cls: type[TaskTransformArgs] 39 | 40 | def __init__( 41 | self, 42 | transform_args: TaskTransformArgs, 43 | ) -> None: 44 | if not isinstance(transform_args, self.transform_args_cls): 45 | raise TypeError( 46 | f"transform_args must be of type {self.transform_args_cls.__name__}, " 47 | f"got {type(transform_args).__name__} instead." 48 | ) 49 | self.transform_args = transform_args 50 | 51 | def __call__(self, input: Any) -> Any: 52 | raise NotImplementedError() 53 | -------------------------------------------------------------------------------- /tests/_methods/simclr/test_simclr.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Literal 11 | 12 | import pytest 13 | 14 | from lightly_train._methods.simclr.simclr import SimCLR, SimCLRArgs, SimCLRSGDArgs 15 | from lightly_train._optim.adamw_args import AdamWArgs 16 | from lightly_train._optim.optimizer_args import OptimizerArgs 17 | from lightly_train._optim.optimizer_type import OptimizerType 18 | from lightly_train._scaling import ScalingInfo 19 | 20 | from ...helpers import DummyCustomModel 21 | 22 | 23 | class TestSimCLRArgs: 24 | def test_resolve_auto(self) -> None: 25 | args = SimCLRArgs() 26 | scaling_info = ScalingInfo(dataset_size=20_000, epochs=100) 27 | args.resolve_auto( 28 | scaling_info=scaling_info, 29 | optimizer_args=AdamWArgs(), 30 | wrapped_model=DummyCustomModel(), 31 | ) 32 | assert not args.has_auto() 33 | 34 | @pytest.mark.parametrize( 35 | "optim_type, expected", 36 | [ 37 | ("auto", SimCLRSGDArgs), 38 | (OptimizerType.ADAMW, AdamWArgs), 39 | (OptimizerType.SGD, SimCLRSGDArgs), 40 | ], 41 | ) 42 | def test_optimizer_args_cls( 43 | self, optim_type: OptimizerType | Literal["auto"], expected: type[OptimizerArgs] 44 | ) -> None: 45 | assert SimCLR.optimizer_args_cls(optim_type=optim_type) == expected 46 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_src/utils/custom_callable.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This software may be used and distributed in accordance with 5 | # the terms of the DINOv3 License Agreement.# 6 | 7 | from __future__ import annotations 8 | 9 | import contextlib 10 | import importlib 11 | import inspect 12 | import os 13 | import sys 14 | from pathlib import Path 15 | 16 | 17 | @contextlib.contextmanager 18 | def _load_modules_from_dir(dir_: str): 19 | sys.path.insert(0, dir_) 20 | yield 21 | sys.path.pop(0) 22 | 23 | 24 | def load_custom_callable(module_path: str | Path, callable_name: str): 25 | module_full_path = os.path.realpath(module_path) 26 | assert os.path.exists(module_full_path), f"module {module_full_path} does not exist" 27 | module_dir, module_filename = os.path.split(module_full_path) 28 | module_name, _ = os.path.splitext(module_filename) 29 | 30 | with _load_modules_from_dir(module_dir): 31 | module = importlib.import_module(module_name) 32 | if inspect.getfile(module) != module_full_path: 33 | importlib.reload(module) 34 | callable_ = getattr(module, callable_name) 35 | 36 | return callable_ 37 | 38 | 39 | @contextlib.contextmanager 40 | def change_working_dir_and_pythonpath(new_dir): 41 | old_dir = Path.cwd() 42 | new_dir = Path(new_dir).expanduser().resolve().as_posix() 43 | old_pythonpath = sys.path.copy() 44 | sys.path.insert(0, new_dir) 45 | os.chdir(new_dir) 46 | try: 47 | yield 48 | finally: 49 | os.chdir(old_dir) 50 | sys.path = old_pythonpath 51 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/dinov2/dinov2_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | from __future__ import annotations 10 | 11 | from pydantic import Field 12 | 13 | from lightly_train._methods.dino.dino_transform import ( 14 | DINOLocalViewRandomResizeArgs, 15 | DINOLocalViewTransformArgs, 16 | DINORandomResizeArgs, 17 | DINOTransform, 18 | DINOTransformArgs, 19 | ) 20 | from lightly_train.types import ImageSizeTuple 21 | 22 | 23 | class DINOv2RandomResizeArgs(DINORandomResizeArgs): 24 | min_scale: float = 0.32 25 | 26 | 27 | class DINOv2LocalViewRandomResizeArgs(DINOLocalViewRandomResizeArgs): 28 | max_scale: float = 0.32 29 | 30 | 31 | class DINOv2ViTLocalViewTransformArgs(DINOLocalViewTransformArgs): 32 | num_views: int = 8 33 | view_size: ImageSizeTuple = (98, 98) 34 | random_resize: DINOv2LocalViewRandomResizeArgs | None = Field( 35 | default_factory=DINOv2LocalViewRandomResizeArgs 36 | ) 37 | 38 | 39 | class DINOv2ViTTransformArgs(DINOTransformArgs): 40 | random_resize: DINOv2RandomResizeArgs | None = Field( 41 | default_factory=DINOv2RandomResizeArgs 42 | ) 43 | local_view: DINOv2ViTLocalViewTransformArgs | None = Field( 44 | default_factory=DINOv2ViTLocalViewTransformArgs 45 | ) 46 | 47 | 48 | class DINOv2ViTTransform(DINOTransform): 49 | @staticmethod 50 | def transform_args_cls() -> type[DINOv2ViTTransformArgs]: 51 | return DINOv2ViTTransformArgs 52 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | ### Documentation 2 | 3 | VERSION ?= $(shell grep '__version__' ../src/lightly_train/__init__.py | sed -E 's/[^0-9.]//g') 4 | BUILD_DIR := build 5 | # The build directory for the current version. We create a new directory for each 6 | # version to simplify hosting the documentation for multiple versions. This follows 7 | # the pattern used by PyTorch: https://pytorch.org/docs/versions.html 8 | BUILD_VERSION_DIR := ${BUILD_DIR}/${VERSION} 9 | BUILD_STABLE_DIR := ${BUILD_DIR}/stable 10 | SOURCE_DIR := source 11 | 12 | 13 | # Build docs for the current version. 14 | .PHONY: docs 15 | docs: 16 | @echo "📚 Building v${VERSION} documentation..." 17 | python prebuild.py --source-dir ${SOURCE_DIR} 18 | sphinx-build -b html --fail-on-warning --keep-going ${SOURCE_DIR} ${BUILD_VERSION_DIR} 19 | python build.py --build-dir ${BUILD_DIR} 20 | @echo "✅ Documentation built successfully!" 21 | 22 | # Build docs for the stable version. Assumes that the current version is the stable 23 | # version. 24 | .PHONY: docs-stable 25 | docs-stable: 26 | @echo "🗿 Building stable documentation..." 27 | python prebuild.py --source-dir ${SOURCE_DIR} 28 | sphinx-build -b html --fail-on-warning --keep-going ${SOURCE_DIR} ${BUILD_STABLE_DIR} 29 | python build.py --build-dir ${BUILD_DIR} 30 | @echo "✅ Documentation built successfully!" 31 | 32 | # Serve the documentation on localhost. 33 | .PHONY: serve 34 | serve: 35 | python -m http.server 1234 --directory ${BUILD_DIR} 36 | 37 | .PHONY: clean 38 | clean: 39 | rm -rf ${BUILD_DIR} 40 | 41 | .PHONY: format 42 | format: 43 | make -C .. format 44 | 45 | .PHONY: format-check 46 | format-check: 47 | make -C .. format-check -------------------------------------------------------------------------------- /src/lightly_train/_models/torchvision/convnext.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from torch import Tensor 9 | from torchvision.models import ConvNeXt 10 | 11 | from lightly_train._models.model_wrapper import ( 12 | ForwardFeaturesOutput, 13 | ForwardPoolOutput, 14 | ) 15 | from lightly_train._models.torchvision.torchvision import TorchvisionModelWrapper 16 | 17 | 18 | class ConvNeXtModelWrapper(TorchvisionModelWrapper): 19 | _torchvision_models = [ConvNeXt] 20 | _torchvision_model_name_pattern = r"convnext.*" 21 | 22 | def __init__(self, model: ConvNeXt) -> None: 23 | super().__init__() 24 | self._model = [model] 25 | self._features = model.features 26 | self._pool = model.avgpool 27 | # Use linear layer from classifier to get feature dimension as last layer of 28 | # `model.features` is different depending on model configuration, making it hard 29 | # to get the feature dimension from there. 30 | self._feature_dim: int = model.classifier[-1].in_features 31 | 32 | def feature_dim(self) -> int: 33 | return self._feature_dim 34 | 35 | def forward_features(self, x: Tensor) -> ForwardFeaturesOutput: 36 | return {"features": self._features(x)} 37 | 38 | def forward_pool(self, x: ForwardFeaturesOutput) -> ForwardPoolOutput: 39 | return {"pooled_features": self._pool(x["features"])} 40 | 41 | def get_model(self) -> ConvNeXt: 42 | return self._model[0] 43 | -------------------------------------------------------------------------------- /src/lightly_train/_models/torchvision/shufflenet.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | from torch import Tensor 10 | from torch.nn import Module 11 | from torchvision.models import ShuffleNetV2 12 | 13 | from lightly_train._models.model_wrapper import ForwardFeaturesOutput, ForwardPoolOutput 14 | from lightly_train._models.torchvision.torchvision import TorchvisionModelWrapper 15 | 16 | 17 | class ShuffleNetV2ModelWrapper(TorchvisionModelWrapper): 18 | _torchvision_models = [ShuffleNetV2] 19 | _torchvision_model_name_pattern = r"shufflenet_v2.*" 20 | 21 | def __init__(self, model: Module): 22 | super().__init__() 23 | self._model = model 24 | 25 | def get_model(self) -> Module: 26 | return self._model 27 | 28 | def forward_features(self, x: Tensor) -> ForwardFeaturesOutput: 29 | x = self._model.conv1(x) # type: ignore 30 | x = self._model.maxpool(x) # type: ignore 31 | x = self._model.stage2(x) # type: ignore 32 | x = self._model.stage3(x) # type: ignore 33 | x = self._model.stage4(x) # type: ignore 34 | x = self._model.conv5(x) # type: ignore 35 | return {"features": x} 36 | 37 | def forward_pool(self, x: ForwardFeaturesOutput) -> ForwardPoolOutput: 38 | return {"pooled_features": x["features"].mean([2, 3], keepdim=True)} 39 | 40 | def feature_dim(self) -> int: 41 | feature_dim: int = self._model.fc.in_features # type: ignore 42 | return feature_dim 43 | -------------------------------------------------------------------------------- /docs/source/python_api/lightly_train.md: -------------------------------------------------------------------------------- 1 | (lightly-train)= 2 | 3 | # lightly_train 4 | 5 | Documentation of the public API of the `lightly_train` package. 6 | 7 | ## Functions 8 | 9 | ```{eval-rst} 10 | 11 | .. automodule:: lightly_train 12 | :members: embed, export, export_onnx, list_methods, list_models, load_model, train, train_instance_segmentation, train_object_detection, train_panoptic_segmentation, train_semantic_segmentation 13 | 14 | ``` 15 | 16 | ## Models 17 | 18 | ```{eval-rst} 19 | 20 | .. autoclass:: lightly_train._task_models.dinov3_eomt_instance_segmentation.task_model.DINOv3EoMTInstanceSegmentation 21 | :members: predict 22 | :exclude-members: __init__, __new__ 23 | 24 | .. autoclass:: lightly_train._task_models.dinov2_ltdetr_object_detection.task_model.DINOv2LTDETRObjectDetection 25 | :members: predict 26 | :exclude-members: __init__, __new__ 27 | 28 | .. autoclass:: lightly_train._task_models.dinov3_ltdetr_object_detection.task_model.DINOv3LTDETRObjectDetection 29 | :members: predict 30 | :exclude-members: __init__, __new__ 31 | 32 | .. autoclass:: lightly_train._task_models.dinov3_eomt_panoptic_segmentation.task_model.DINOv3EoMTPanopticSegmentation 33 | :members: predict 34 | :exclude-members: __init__, __new__ 35 | 36 | .. autoclass:: lightly_train._task_models.dinov2_eomt_semantic_segmentation.task_model.DINOv2EoMTSemanticSegmentation 37 | :members: predict 38 | :exclude-members: __init__, __new__ 39 | 40 | .. autoclass:: lightly_train._task_models.dinov3_eomt_semantic_segmentation.task_model.DINOv3EoMTSemanticSegmentation 41 | :members: predict 42 | :exclude-members: __init__, __new__ 43 | 44 | ``` 45 | -------------------------------------------------------------------------------- /tests/_methods/densecl/test_densecl.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Literal 11 | 12 | import pytest 13 | 14 | from lightly_train._methods.densecl.densecl import DenseCL, DenseCLArgs, DenseCLSGDArgs 15 | from lightly_train._optim.adamw_args import AdamWArgs 16 | from lightly_train._optim.optimizer_args import OptimizerArgs 17 | from lightly_train._optim.optimizer_type import OptimizerType 18 | from lightly_train._scaling import ScalingInfo 19 | 20 | from ...helpers import DummyCustomModel 21 | 22 | 23 | class TestDenseCLArgs: 24 | def test_resolve_auto(self) -> None: 25 | args = DenseCLArgs() 26 | scaling_info = ScalingInfo(dataset_size=20_000, epochs=100) 27 | args.resolve_auto( 28 | scaling_info=scaling_info, 29 | optimizer_args=AdamWArgs(), 30 | wrapped_model=DummyCustomModel(), 31 | ) 32 | assert args.memory_bank_size == 8192 33 | assert not args.has_auto() 34 | 35 | @pytest.mark.parametrize( 36 | "optim_type, expected", 37 | [ 38 | ("auto", DenseCLSGDArgs), 39 | (OptimizerType.ADAMW, AdamWArgs), 40 | (OptimizerType.SGD, DenseCLSGDArgs), 41 | ], 42 | ) 43 | def test_optimizer_args_cls( 44 | self, optim_type: OptimizerType | Literal["auto"], expected: type[OptimizerArgs] 45 | ) -> None: 46 | assert DenseCL.optimizer_args_cls(optim_type=optim_type) == expected 47 | -------------------------------------------------------------------------------- /.github/workflows/test_build.yml: -------------------------------------------------------------------------------- 1 | name: Test Build 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | workflow_dispatch: 9 | 10 | env: 11 | # Install packages into system environment. 12 | # Follows: https://docs.astral.sh/uv/guides/integration/github/#using-uv-pip 13 | UV_SYSTEM_PYTHON: 1 14 | 15 | jobs: 16 | detect-code-changes: 17 | name: Detect Code Changes 18 | runs-on: ubuntu-latest 19 | outputs: 20 | run-tests: ${{ steps.filter.outputs.run-tests }} 21 | steps: 22 | - uses: actions/checkout@v4 23 | - uses: dorny/paths-filter@v3 24 | id: filter 25 | with: 26 | list-files: shell 27 | filters: | 28 | run-tests: 29 | - '!docs/**' 30 | - '!docker/**' 31 | - '!.github/**' 32 | - '.github/workflows/test_build.yml' 33 | test-build: 34 | name: Test Build 35 | needs: detect-code-changes 36 | if: needs.detect-code-changes.outputs.run-tests == 'true' 37 | runs-on: ubuntu-latest 38 | steps: 39 | - name: Checkout Code 40 | uses: actions/checkout@v4 41 | - name: Set Up uv 42 | uses: astral-sh/setup-uv@v5 43 | id: setup-uv 44 | with: 45 | version: 0.6.11 46 | enable-cache: true 47 | cache-dependency-glob: "**/pyproject.toml" 48 | python-version: "3.10" 49 | - name: Set Up Python 50 | uses: actions/setup-python@v5 51 | with: 52 | python-version: "3.10" 53 | - name: Set Up Environment 54 | run: | 55 | make install-dist 56 | - name: Build 57 | run: | 58 | make dist 59 | -------------------------------------------------------------------------------- /src/lightly_train/_embedding/writers/torch_writer.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pathlib import Path 11 | 12 | import torch 13 | from torch import Tensor 14 | 15 | from lightly_train._embedding.embedding_format import EmbeddingFormat 16 | from lightly_train._embedding.writers.embedding_writer import EmbeddingWriter 17 | 18 | 19 | class TorchWriter(EmbeddingWriter): 20 | """Writes embeddings to a file in torch format. 21 | 22 | Args: 23 | filepath: 24 | Path to file where embeddings will be saved. 25 | 26 | Example output: 27 | ``` 28 | embeddings = torch.load("embeddings.pt") 29 | embeddings == { 30 | "filenames": ["image1.jpg", "image2.jpg", ...], 31 | "embeddings": torch.tensor([ 32 | [0.1, 0.2, 0.3], 33 | [0.4, 0.5, 0.6], 34 | ... 35 | ]) 36 | } 37 | ``` 38 | 39 | """ 40 | 41 | def __init__(self, filepath: Path): 42 | super().__init__() 43 | self._filepath = filepath 44 | 45 | @classmethod 46 | def is_supported_format(cls, format: EmbeddingFormat) -> bool: 47 | return format == EmbeddingFormat.TORCH 48 | 49 | def save(self, embeddings: Tensor | None, filenames: list[str]) -> None: 50 | torch.save( 51 | { 52 | "filenames": filenames, 53 | "embeddings": torch.empty(0) if embeddings is None else embeddings, 54 | }, 55 | self._filepath, 56 | ) 57 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/distillationv2/distillationv2_loss.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from torch import Tensor 11 | from torch.nn import Module, MSELoss 12 | 13 | 14 | class DistillationV2Loss(Module): 15 | """ 16 | Computes the Mean Squared Error (MSE) loss used for DistillationV2. 17 | 18 | The loss directly compares the student and teacher representations using the 19 | MSE loss function. The student and teacher features are not normalized. 20 | 21 | """ 22 | 23 | def __init__( 24 | self, 25 | ) -> None: 26 | super().__init__() 27 | self.mse_loss = MSELoss() 28 | 29 | def forward(self, teacher_features: Tensor, student_features: Tensor) -> Tensor: 30 | """Computes the MSE loss between the student and teacher features. 31 | 32 | Args: 33 | teacher_features: Tensor containing teacher representations from the current batch. 34 | The expected shape is (batch_size, n_features, feature_dim). 35 | n_features is the number of tokens per sequence in the teacher model. 36 | student_features: Tensor containing student representations from the current batch. 37 | The expected shape is (batch_size, n_features, feature_dim). 38 | 39 | Returns: 40 | MSE loss as a scalar tensor. 41 | """ 42 | # Compute the loss. 43 | loss: Tensor = self.mse_loss(teacher_features, student_features) 44 | return loss 45 | -------------------------------------------------------------------------------- /tests/_transforms/test_normalize.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import numpy as np 11 | 12 | from lightly_train._transforms.normalize import NormalizeDtypeAware 13 | 14 | 15 | def test_standard_normalization_float_image() -> None: 16 | transform = NormalizeDtypeAware() 17 | image = np.array( 18 | [ 19 | [[0.0, 0.25, 0.5], [0.75, 1.0, 0.5]], 20 | [[0.1, 0.2, 0.3], [0.9, 0.8, 0.7]], 21 | ], 22 | dtype=np.float32, 23 | ) 24 | 25 | result = transform(image=image)["image"] 26 | 27 | mean = np.asarray(transform.mean, dtype=np.float32) 28 | std = np.asarray(transform.std, dtype=np.float32) 29 | 30 | expected = (image.astype(np.float32) - mean) / std 31 | np.testing.assert_allclose(result, expected, atol=1e-6) 32 | 33 | 34 | def test_standard_normalization_uint8_image() -> None: 35 | transform = NormalizeDtypeAware() 36 | image = np.array( 37 | [ 38 | [[0, 64, 128], [192, 255, 32]], 39 | [[16, 48, 80], [112, 144, 176]], 40 | ], 41 | dtype=np.uint8, 42 | ) 43 | 44 | result = transform(image=image)["image"] 45 | 46 | mean = np.asarray(transform.mean, dtype=np.float32) * float( 47 | transform.max_pixel_value 48 | ) 49 | std = np.asarray(transform.std, dtype=np.float32) * float(transform.max_pixel_value) 50 | 51 | expected = (image.astype(np.float32) - mean) / std 52 | np.testing.assert_allclose(result, expected, atol=1e-6) 53 | -------------------------------------------------------------------------------- /.github/workflows/test_documentation.yml: -------------------------------------------------------------------------------- 1 | name: Test Documentation 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | workflow_dispatch: 9 | 10 | env: 11 | # Install packages into system environment. 12 | # Follows: https://docs.astral.sh/uv/guides/integration/github/#using-uv-pip 13 | UV_SYSTEM_PYTHON: 1 14 | 15 | jobs: 16 | detect-code-changes: 17 | name: Detect Code Changes 18 | runs-on: ubuntu-latest 19 | outputs: 20 | run-tests: ${{ steps.filter.outputs.run-tests }} 21 | steps: 22 | - uses: actions/checkout@v4 23 | - uses: dorny/paths-filter@v3 24 | id: filter 25 | with: 26 | list-files: shell 27 | filters: | 28 | run-tests: 29 | - '!docker/**' 30 | - '!.github/**' 31 | - '.github/workflows/test_documentation.yml' 32 | test-documentation: 33 | name: Test Documentation 34 | needs: detect-code-changes 35 | if: needs.detect-code-changes.outputs.run-tests == 'true' 36 | runs-on: ubuntu-latest 37 | steps: 38 | - name: Checkout code 39 | uses: actions/checkout@v4 40 | - name: Set Up uv 41 | uses: astral-sh/setup-uv@v5 42 | id: setup-uv 43 | with: 44 | version: 0.6.11 45 | enable-cache: true 46 | cache-dependency-glob: "**/pyproject.toml" 47 | python-version: "3.10" 48 | - name: Set Up Python 49 | uses: actions/setup-python@v5 50 | with: 51 | python-version: "3.10" 52 | - name: Set Up Environment 53 | run: | 54 | make install-docs 55 | - name: Build 56 | run: | 57 | cd docs && make docs docs-stable 58 | -------------------------------------------------------------------------------- /tests/_models/test_model_wrapper.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import pytest 11 | 12 | from lightly_train._models import model_wrapper 13 | 14 | from ..helpers import DummyCustomModel 15 | 16 | 17 | @pytest.mark.parametrize( 18 | "exclude_module_attrs", 19 | [False, True], 20 | ) 21 | def test_missing_model_wrapper_attrs(exclude_module_attrs: bool) -> None: 22 | assert not model_wrapper.missing_model_wrapper_attrs( 23 | DummyCustomModel(), exclude_module_attrs=exclude_module_attrs 24 | ) 25 | 26 | 27 | @pytest.mark.parametrize( 28 | "exclude_module_attrs, expected", 29 | [ 30 | (True, ["feature_dim", "forward_pool"]), 31 | ( 32 | False, 33 | [ 34 | "T_destination", 35 | "feature_dim", 36 | "forward_pool", 37 | "load_state_dict", 38 | "parameters", 39 | "state_dict", 40 | ], 41 | ), 42 | ], 43 | ) 44 | def test_missing_model_wrapper_attrs__missing( 45 | exclude_module_attrs: bool, expected: list[str] 46 | ) -> None: 47 | class InvalidCustomModelWrapper: 48 | def get_model(self) -> None: 49 | pass 50 | 51 | def forward_features(self) -> None: 52 | pass 53 | 54 | assert ( 55 | model_wrapper.missing_model_wrapper_attrs( 56 | InvalidCustomModelWrapper(), 57 | exclude_module_attrs=exclude_module_attrs, 58 | ) 59 | == expected 60 | ) 61 | -------------------------------------------------------------------------------- /tests/test__scaling.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import pytest 11 | 12 | from lightly_train import _scaling 13 | 14 | 15 | @pytest.mark.parametrize( 16 | "input, input_start, input_end, value_start, value_end, expected", 17 | [ 18 | (0, 0, 1, 10, 20, 10), 19 | (0.5, 0, 1, 10, 20, 15), 20 | (1, 0, 1, 10, 20, 20), 21 | (-1, 0, 1, 10, 20, 10), # clamp to lower value 22 | (2, 0, 1, 10, 20, 20), # clamp to upper value 23 | ], 24 | ) 25 | def test_interpolate( 26 | input: float, 27 | input_start: float, 28 | input_end: float, 29 | value_start: float, 30 | value_end: float, 31 | expected: float, 32 | ) -> None: 33 | assert ( 34 | _scaling.interpolate( 35 | input=input, 36 | input_start=input_start, 37 | input_end=input_end, 38 | value_start=value_start, 39 | value_end=value_end, 40 | ) 41 | == expected 42 | ) 43 | 44 | 45 | @pytest.mark.parametrize( 46 | "input, buckets, expected", 47 | [ 48 | (0, [(1, 10), (2, 20)], 10), 49 | (1, [(1, 10), (2, 20)], 20), 50 | ], 51 | ) 52 | def test_get_bucket_value( 53 | input: int, buckets: list[tuple[int, int]], expected: int 54 | ) -> None: 55 | assert _scaling.get_bucket_value(input=input, buckets=buckets) == expected 56 | 57 | 58 | def test_get_bucket_value__input_too_large() -> None: 59 | with pytest.raises(ValueError): 60 | _scaling.get_bucket_value(input=2, buckets=[(1, 10), (2, 20)]) 61 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov2_eomt_semantic_segmentation/scale_block.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | # Modifications Copyright 2025 Lightly AG: 9 | # - Replace timm LayerNorm2D implementation with torch version 10 | 11 | from __future__ import annotations 12 | 13 | from typing import Type 14 | 15 | from torch import Tensor, nn 16 | from torch.nn import LayerNorm, Module 17 | 18 | 19 | class ScaleBlock(Module): 20 | def __init__(self, embed_dim: int, conv1_layer: Type[Module] = nn.ConvTranspose2d): 21 | super().__init__() 22 | 23 | self.conv1 = conv1_layer( 24 | embed_dim, 25 | embed_dim, 26 | kernel_size=2, 27 | stride=2, 28 | ) 29 | self.act = nn.GELU() 30 | self.conv2 = nn.Conv2d( 31 | embed_dim, 32 | embed_dim, 33 | kernel_size=3, 34 | padding=1, 35 | groups=embed_dim, 36 | bias=False, 37 | ) 38 | self.norm = LayerNorm2D(embed_dim) 39 | 40 | def forward(self, x: Tensor) -> Tensor: 41 | x = self.conv1(x) 42 | x = self.act(x) 43 | x = self.conv2(x) 44 | x = self.norm(x) 45 | return x 46 | 47 | 48 | class LayerNorm2D(LayerNorm): 49 | def __init__(self, embed_dim: int): 50 | super().__init__(normalized_shape=embed_dim) 51 | 52 | def forward(self, x: Tensor) -> Tensor: 53 | x = x.permute(0, 2, 3, 1) 54 | x = super().forward(x) 55 | x = x.permute(0, 3, 1, 2) 56 | return x 57 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_instance_segmentation/scale_block.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | # Modifications Copyright 2025 Lightly AG: 9 | # - Replace timm LayerNorm2D implementation with torch version 10 | 11 | from __future__ import annotations 12 | 13 | from typing import Type 14 | 15 | from torch import Tensor, nn 16 | from torch.nn import LayerNorm, Module 17 | 18 | 19 | class ScaleBlock(Module): 20 | def __init__(self, embed_dim: int, conv1_layer: Type[Module] = nn.ConvTranspose2d): 21 | super().__init__() 22 | 23 | self.conv1 = conv1_layer( 24 | embed_dim, 25 | embed_dim, 26 | kernel_size=2, 27 | stride=2, 28 | ) 29 | self.act = nn.GELU() 30 | self.conv2 = nn.Conv2d( 31 | embed_dim, 32 | embed_dim, 33 | kernel_size=3, 34 | padding=1, 35 | groups=embed_dim, 36 | bias=False, 37 | ) 38 | self.norm = LayerNorm2D(embed_dim) 39 | 40 | def forward(self, x: Tensor) -> Tensor: 41 | x = self.conv1(x) 42 | x = self.act(x) 43 | x = self.conv2(x) 44 | x = self.norm(x) 45 | return x 46 | 47 | 48 | class LayerNorm2D(LayerNorm): 49 | def __init__(self, embed_dim: int): 50 | super().__init__(normalized_shape=embed_dim) 51 | 52 | def forward(self, x: Tensor) -> Tensor: 53 | x = x.permute(0, 2, 3, 1) 54 | x = super().forward(x) 55 | x = x.permute(0, 3, 1, 2) 56 | return x 57 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_panoptic_segmentation/scale_block.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | # Modifications Copyright 2025 Lightly AG: 9 | # - Replace timm LayerNorm2D implementation with torch version 10 | 11 | from __future__ import annotations 12 | 13 | from typing import Type 14 | 15 | from torch import Tensor, nn 16 | from torch.nn import LayerNorm, Module 17 | 18 | 19 | class ScaleBlock(Module): 20 | def __init__(self, embed_dim: int, conv1_layer: Type[Module] = nn.ConvTranspose2d): 21 | super().__init__() 22 | 23 | self.conv1 = conv1_layer( 24 | embed_dim, 25 | embed_dim, 26 | kernel_size=2, 27 | stride=2, 28 | ) 29 | self.act = nn.GELU() 30 | self.conv2 = nn.Conv2d( 31 | embed_dim, 32 | embed_dim, 33 | kernel_size=3, 34 | padding=1, 35 | groups=embed_dim, 36 | bias=False, 37 | ) 38 | self.norm = LayerNorm2D(embed_dim) 39 | 40 | def forward(self, x: Tensor) -> Tensor: 41 | x = self.conv1(x) 42 | x = self.act(x) 43 | x = self.conv2(x) 44 | x = self.norm(x) 45 | return x 46 | 47 | 48 | class LayerNorm2D(LayerNorm): 49 | def __init__(self, embed_dim: int): 50 | super().__init__(normalized_shape=embed_dim) 51 | 52 | def forward(self, x: Tensor) -> Tensor: 53 | x = x.permute(0, 2, 3, 1) 54 | x = super().forward(x) 55 | x = x.permute(0, 3, 1, 2) 56 | return x 57 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_semantic_segmentation/scale_block.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | # Modifications Copyright 2025 Lightly AG: 9 | # - Replace timm LayerNorm2D implementation with torch version 10 | 11 | from __future__ import annotations 12 | 13 | from typing import Type 14 | 15 | from torch import Tensor, nn 16 | from torch.nn import LayerNorm, Module 17 | 18 | 19 | class ScaleBlock(Module): 20 | def __init__(self, embed_dim: int, conv1_layer: Type[Module] = nn.ConvTranspose2d): 21 | super().__init__() 22 | 23 | self.conv1 = conv1_layer( 24 | embed_dim, 25 | embed_dim, 26 | kernel_size=2, 27 | stride=2, 28 | ) 29 | self.act = nn.GELU() 30 | self.conv2 = nn.Conv2d( 31 | embed_dim, 32 | embed_dim, 33 | kernel_size=3, 34 | padding=1, 35 | groups=embed_dim, 36 | bias=False, 37 | ) 38 | self.norm = LayerNorm2D(embed_dim) 39 | 40 | def forward(self, x: Tensor) -> Tensor: 41 | x = self.conv1(x) 42 | x = self.act(x) 43 | x = self.conv2(x) 44 | x = self.norm(x) 45 | return x 46 | 47 | 48 | class LayerNorm2D(LayerNorm): 49 | def __init__(self, embed_dim: int): 50 | super().__init__(normalized_shape=embed_dim) 51 | 52 | def forward(self, x: Tensor) -> Tensor: 53 | x = x.permute(0, 2, 3, 1) 54 | x = super().forward(x) 55 | x = x.permute(0, 3, 1, 2) 56 | return x 57 | -------------------------------------------------------------------------------- /src/lightly_train/_data/task_dataset.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from collections.abc import Iterable, Sequence 11 | from typing import ClassVar 12 | 13 | from torch.utils.data import Dataset 14 | 15 | from lightly_train._configs.config import PydanticConfig 16 | from lightly_train._data.task_batch_collation import BaseCollateFunction 17 | from lightly_train._transforms.task_transform import TaskTransform 18 | from lightly_train.types import TaskDatasetItem 19 | 20 | 21 | class TaskDatasetArgs(PydanticConfig): 22 | def list_image_info(self) -> Iterable[dict[str, str]]: 23 | """Listing the image info should not happen in-memory for large datasets.""" 24 | raise NotImplementedError() 25 | 26 | def get_dataset_cls(self) -> type[TaskDataset]: 27 | raise NotImplementedError() 28 | 29 | 30 | class TaskDataset(Dataset[TaskDatasetItem]): 31 | batch_collate_fn_cls: ClassVar[type[BaseCollateFunction]] = BaseCollateFunction 32 | 33 | def __init__( 34 | self, 35 | dataset_args: TaskDatasetArgs, 36 | image_info: Sequence[dict[str, str]], 37 | transform: TaskTransform, 38 | ) -> None: 39 | self.dataset_args = dataset_args 40 | self.image_info = image_info 41 | self._transform = transform 42 | 43 | @property 44 | def transform(self) -> TaskTransform: 45 | return self._transform 46 | 47 | def __len__(self) -> int: 48 | return len(self.image_info) 49 | 50 | def __getitem__(self, index: int) -> TaskDatasetItem: 51 | raise NotImplementedError() 52 | -------------------------------------------------------------------------------- /tests/_embedding/writers/test_torch_writer.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pathlib import Path 11 | 12 | import pytest 13 | import torch 14 | 15 | from lightly_train._embedding.embedding_format import EmbeddingFormat 16 | from lightly_train._embedding.writers.torch_writer import TorchWriter 17 | 18 | 19 | class TestTorchWriter: 20 | @pytest.mark.parametrize( 21 | "format, expected", 22 | [ 23 | (EmbeddingFormat.TORCH, True), 24 | (EmbeddingFormat.CSV, False), 25 | ], 26 | ) 27 | def test_is_supported_format(self, format: EmbeddingFormat, expected: bool) -> None: 28 | assert TorchWriter.is_supported_format(format=format) is expected 29 | 30 | def test_save(self, tmp_path: Path) -> None: 31 | filepath = tmp_path / "embeddings.pt" 32 | writer = TorchWriter(filepath=filepath) 33 | embeddings = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) 34 | filenames = ["file1.png", "file , special.png"] 35 | writer.save(embeddings=embeddings, filenames=filenames) 36 | 37 | loaded = torch.load(filepath) 38 | assert loaded["filenames"] == filenames 39 | assert torch.equal(loaded["embeddings"], embeddings) 40 | 41 | def test_save__none(self, tmp_path: Path) -> None: 42 | filepath = tmp_path / "embeddings.pt" 43 | writer = TorchWriter(filepath=filepath) 44 | writer.save(embeddings=None, filenames=[]) 45 | loaded = torch.load(filepath) 46 | assert loaded["filenames"] == [] 47 | assert loaded["embeddings"].shape == (0,) 48 | -------------------------------------------------------------------------------- /tests/_configs/test_config.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from typing import Any 11 | 12 | import pytest 13 | 14 | from lightly_train._configs.config import PydanticConfig 15 | 16 | 17 | class _Config(PydanticConfig): 18 | a: Any 19 | 20 | 21 | class _Default(PydanticConfig): 22 | a: str = "auto" 23 | 24 | 25 | class TestPydanticConfig: 26 | @pytest.mark.parametrize( 27 | "config, expected", 28 | [ 29 | # Config 30 | (_Config(a=""), False), 31 | (_Config(a="test"), False), 32 | (_Config(a="auto"), True), 33 | # Dict 34 | (_Config(a={"a": ""}), False), 35 | (_Config(a={"a": "test"}), False), 36 | (_Config(a={"a": "auto"}), True), 37 | # Config in config 38 | (_Config(a=_Config(a="")), False), 39 | (_Config(a=_Config(a="auto")), True), 40 | # Config in dict 41 | (_Config(a={"a": _Config(a="")}), False), 42 | (_Config(a={"a": _Config(a="auto")}), True), 43 | # Dict in config 44 | (_Config(a=_Config(a={"a": ""})), False), 45 | (_Config(a=_Config(a={"a": "auto"})), True), 46 | # Dict in dict 47 | (_Config(a={"a": {"b": "auto"}}), True), 48 | # Default 49 | (_Default(), True), 50 | (_Default(a=""), False), 51 | (_Config(a=_Default()), True), 52 | (_Config(a=_Default(a="")), False), 53 | ], 54 | ) 55 | def test_has_auto(self, config: PydanticConfig, expected: bool) -> None: 56 | assert config.has_auto() == expected 57 | -------------------------------------------------------------------------------- /tests/_data/test_cache.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import os 9 | from pathlib import Path 10 | 11 | from pytest_mock import MockerFixture 12 | 13 | from lightly_train._commands import common_helpers 14 | from lightly_train._data import cache 15 | 16 | 17 | def test_get_model_cache_dir__default(mocker: MockerFixture) -> None: 18 | mocker.patch.dict(os.environ, {"LIGHTLY_TRAIN_CACHE_DIR": ""}) 19 | expected = Path.home() / ".cache" / "lightly-train" / "models" 20 | assert cache.get_model_cache_dir() == expected 21 | 22 | 23 | def test_get_model_cache_dir__custom(tmp_path: Path, mocker: MockerFixture) -> None: 24 | mocker.patch.dict(os.environ, {"LIGHTLY_TRAIN_MODEL_CACHE_DIR": str(tmp_path)}) 25 | assert cache.get_model_cache_dir() == tmp_path 26 | 27 | 28 | def test_get_data_cache_dir__default(mocker: MockerFixture) -> None: 29 | mocker.patch.dict(os.environ, {"LIGHTLY_TRAIN_CACHE_DIR": ""}) 30 | expected = Path.home() / ".cache" / "lightly-train" / "data" 31 | assert cache.get_data_cache_dir() == expected 32 | 33 | 34 | def test_get_data_cache_dir__custom(tmp_path: Path, mocker: MockerFixture) -> None: 35 | mocker.patch.dict(os.environ, {"LIGHTLY_TRAIN_DATA_CACHE_DIR": str(tmp_path)}) 36 | assert cache.get_data_cache_dir() == tmp_path 37 | 38 | 39 | def test_get_data_cache_dir__tmp_dir_propagation( 40 | tmp_path: Path, mocker: MockerFixture 41 | ) -> None: 42 | # Ensure that the temporary directory is propagated to the data cache directory. 43 | mocker.patch.dict(os.environ, {"LIGHTLY_TRAIN_TMP_DIR": str(tmp_path)}) 44 | assert cache.get_data_cache_dir() == tmp_path 45 | assert cache.get_data_cache_dir() == common_helpers.get_tmp_dir() 46 | -------------------------------------------------------------------------------- /tests/test__checkpoint.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from datetime import datetime, timezone 9 | 10 | from lightly_train._checkpoint import ( 11 | CheckpointLightlyTrain, 12 | CheckpointLightlyTrainModels, 13 | ) 14 | from lightly_train._models.embedding_model import EmbeddingModel 15 | from lightly_train._transforms.transform import NormalizeArgs 16 | 17 | from . import helpers 18 | 19 | 20 | class TestCheckpointInfo: 21 | def test_to_dict_from_dict(self) -> None: 22 | date = datetime( 23 | year=2024, month=1, day=2, minute=3, second=4, tzinfo=timezone.utc 24 | ) 25 | wrapped_model = helpers.DummyCustomModel() 26 | embedding_model = EmbeddingModel(wrapped_model=wrapped_model) 27 | info = CheckpointLightlyTrain( 28 | version="abc", 29 | date=date, 30 | models=CheckpointLightlyTrainModels( 31 | model=wrapped_model.get_model(), 32 | wrapped_model=wrapped_model, 33 | embedding_model=embedding_model, 34 | ), 35 | normalize_args=NormalizeArgs(), 36 | ) 37 | 38 | # Check that to_dict representation is correct. 39 | info_dict = info.to_dict() 40 | assert info_dict == { 41 | "version": "abc", 42 | "date": "2024-01-02T00:03:04+00:00", 43 | "models": { 44 | "model": wrapped_model.get_model(), 45 | "wrapped_model": wrapped_model, 46 | "embedding_model": embedding_model, 47 | }, 48 | "normalize_args": NormalizeArgs().to_dict(), 49 | } 50 | 51 | # Check that from_dict reconstruction is correct. 52 | assert info == CheckpointLightlyTrain.from_dict(info_dict) 53 | -------------------------------------------------------------------------------- /src/lightly_train/_methods/densecl/densecl_loss.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import torch 11 | from torch import Tensor 12 | from torch.nn import CrossEntropyLoss, Module 13 | from torch.nn import functional as F 14 | 15 | 16 | class DenseCLLoss(Module): 17 | """DenseCLLoss implementation. 18 | 19 | Uses 'out0' and 'out1' for positives and 'negatives' for negatives. 20 | 21 | TODO(Philipp, 11/24): Move this to LightlySSL once we're certain DenseCL works. 22 | """ 23 | 24 | def __init__(self, temperature: float): 25 | super().__init__() 26 | self.temperature = temperature 27 | self.cross_entropy = CrossEntropyLoss(reduction="mean") 28 | 29 | def forward(self, out0: Tensor, out1: Tensor, negatives: Tensor) -> Tensor: 30 | # Normalize the output to length 1 31 | out0 = F.normalize(out0, dim=-1) 32 | out1 = F.normalize(out1, dim=-1) 33 | 34 | # sim_pos is of shape (batch_size, 1) and sim_pos[i] denotes the similarity 35 | # of the i-th sample in the batch to its positive pair 36 | sim_pos = torch.einsum("nc,nc->n", out0, out1).unsqueeze(-1) 37 | 38 | # sim_neg is of shape (batch_size, memory_bank_size) and sim_neg[i,j] denotes the similarity 39 | # of the i-th sample to the j-th negative sample 40 | sim_neg = torch.einsum("nc,ck->nk", out0, negatives) 41 | 42 | # Set the labels to maximize sim_pos in relation to sim_neg 43 | logits = torch.cat([sim_pos, sim_neg], dim=1) / self.temperature 44 | labels = torch.zeros(logits.shape[0], device=out0.device, dtype=torch.long) 45 | 46 | # Calculate the cross-entropy loss 47 | loss: Tensor = self.cross_entropy(logits, labels) 48 | return loss 49 | -------------------------------------------------------------------------------- /tests/_models/torchvision/test_resnet.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import torch 9 | from torchvision import models 10 | 11 | from lightly_train._models.torchvision.resnet import ResNetModelWrapper 12 | 13 | 14 | class TestResNetModelWrapper: 15 | def test_feature_dim(self) -> None: 16 | model = models.resnet18() 17 | feature_extractor = ResNetModelWrapper(model=model) 18 | assert feature_extractor.feature_dim() == 512 19 | 20 | def test_forward_features(self) -> None: 21 | model = models.resnet18() 22 | feature_extractor = ResNetModelWrapper(model=model) 23 | x = torch.rand(1, 3, 224, 224) 24 | features = feature_extractor.forward_features(x)["features"] 25 | assert features.shape == (1, 512, 7, 7) 26 | 27 | def test_forward_pool(self) -> None: 28 | model = models.resnet18() 29 | feature_extractor = ResNetModelWrapper(model=model) 30 | x = torch.rand(1, 512, 7, 7) 31 | pool = feature_extractor.forward_pool({"features": x})["pooled_features"] 32 | assert pool.shape == (1, 512, 1, 1) 33 | 34 | def test_get_model(self) -> None: 35 | model = models.resnet18() 36 | feature_extractor = ResNetModelWrapper(model=model) 37 | assert feature_extractor.get_model() is model 38 | 39 | def test__device(self) -> None: 40 | # If this test fails it means the wrapped model doesn't move all required 41 | # modules to the correct device. This happens if not all required modules 42 | # are registered as attributes of the class. 43 | model = models.resnet18() 44 | wrapped_model = ResNetModelWrapper(model=model) 45 | wrapped_model.to("meta") 46 | wrapped_model.forward_features(torch.rand(1, 3, 224, 224, device="meta")) 47 | -------------------------------------------------------------------------------- /tests/_methods/test_method_helpers.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import pytest 11 | 12 | from lightly_train._methods import method_helpers 13 | from lightly_train._methods.densecl.densecl import DenseCL 14 | from lightly_train._methods.dino.dino import DINO 15 | from lightly_train._methods.dinov2.dinov2 import DINOv2 16 | from lightly_train._methods.distillation.distillation import Distillation 17 | from lightly_train._methods.distillationv2.distillationv2 import DistillationV2 18 | from lightly_train._methods.method import Method 19 | from lightly_train._methods.simclr.simclr import SimCLR 20 | 21 | from .. import helpers 22 | from ..helpers import DummyCustomModel 23 | 24 | 25 | @pytest.mark.parametrize( 26 | "method, expected", 27 | [ 28 | ("densecl", DenseCL), 29 | ("dino", DINO), 30 | ("dinov2", DINOv2), 31 | ("simclr", SimCLR), 32 | ("distillationv1", Distillation), 33 | ("distillationv2", DistillationV2), 34 | (helpers.get_method(wrapped_model=DummyCustomModel()), SimCLR), 35 | ], 36 | ) 37 | def test_get_method_cls(method: str, expected: type[Method]) -> None: 38 | assert method_helpers.get_method_cls(method=method) == expected 39 | 40 | 41 | def test_list_methods_private() -> None: 42 | assert method_helpers._list_methods() == [ 43 | "densecl", 44 | "dino", 45 | "dinov2", 46 | "distillation", 47 | "distillationv1", 48 | "distillationv2", 49 | "simclr", 50 | ] 51 | 52 | 53 | def test_list_methods_public() -> None: 54 | assert method_helpers.list_methods() == [ 55 | "dino", 56 | "dinov2", 57 | "distillation", 58 | "distillationv1", 59 | "distillationv2", 60 | "simclr", 61 | ] 62 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov2_vit/dinov2_vit_src/models/__init__.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Meta Platforms, Inc. and affiliates. 3 | # 4 | # This source code is licensed under the Apache License, Version 2.0 5 | # found in the LICENSE file in the root directory of this source tree. 6 | # 7 | 8 | 9 | from typing import Tuple, Union 10 | 11 | from torch.nn import Module 12 | 13 | from lightly_train._models.dinov2_vit.dinov2_vit_src.models import ( 14 | vision_transformer as vits, 15 | ) 16 | 17 | 18 | def build_model( 19 | args, only_teacher=False, img_size=224 20 | ) -> Union[Tuple[Module, int], Tuple[Module, Module, int]]: 21 | suffix = "_memeff" 22 | if args.arch.endswith(suffix): 23 | args.arch = args.arch[: -len(suffix)] 24 | if "vit" in args.arch: 25 | vit_kwargs = dict( 26 | img_size=img_size, 27 | patch_size=args.patch_size, 28 | init_values=args.layerscale, 29 | ffn_layer=args.ffn_layer, 30 | block_chunks=args.block_chunks, 31 | qkv_bias=args.qkv_bias, 32 | proj_bias=args.proj_bias, 33 | ffn_bias=args.ffn_bias, 34 | num_register_tokens=args.num_register_tokens, 35 | interpolate_offset=args.interpolate_offset, 36 | interpolate_antialias=args.interpolate_antialias, 37 | ) 38 | teacher = vits.__dict__[args.arch](**vit_kwargs) 39 | if only_teacher: 40 | return teacher, teacher.embed_dim 41 | student = vits.__dict__[args.arch]( 42 | **vit_kwargs, 43 | drop_path_rate=args.drop_path_rate, 44 | drop_path_uniform=args.drop_path_uniform, 45 | ) 46 | embed_dim = student.embed_dim 47 | return student, teacher, embed_dim 48 | 49 | 50 | def build_model_from_cfg(cfg, only_teacher=False) -> Union[Tuple[Module, Module, int],]: 51 | return build_model( 52 | cfg.student, only_teacher=only_teacher, img_size=cfg.crops.global_crops_size 53 | ) 54 | -------------------------------------------------------------------------------- /tests/_models/torchvision/test_torchvision_package.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from pathlib import Path 9 | 10 | import pytest 11 | import torch 12 | from torchvision import models as torchvision_models 13 | 14 | from lightly_train._models.torchvision.torchvision_package import TorchvisionPackage 15 | 16 | from ...helpers import DummyCustomModel 17 | 18 | 19 | class TestTorchvisionPackage: 20 | @pytest.mark.parametrize( 21 | "model_name", 22 | [ 23 | "torchvision/resnet18", 24 | "torchvision/convnext_small", 25 | "torchvision/shufflenet_v2_x0_5", 26 | ], 27 | ) 28 | def test_list_model_names(self, model_name: str) -> None: 29 | assert model_name in TorchvisionPackage.list_model_names() 30 | 31 | def test_is_supported_model__true(self) -> None: 32 | model = torchvision_models.resnet18() 33 | assert TorchvisionPackage.is_supported_model(model) 34 | 35 | wrapped_model = TorchvisionPackage.get_model_wrapper(model=model) 36 | assert TorchvisionPackage.is_supported_model(wrapped_model) 37 | 38 | def test_is_supported_model__false(self) -> None: 39 | model = DummyCustomModel() 40 | assert not TorchvisionPackage.is_supported_model(model=model) 41 | assert not TorchvisionPackage.is_supported_model(model=model.get_model()) 42 | 43 | def test_export_model(self, tmp_path: Path) -> None: 44 | model = torchvision_models.resnet18() 45 | out_path = tmp_path / "model.pt" 46 | TorchvisionPackage.export_model(model=model, out=out_path) 47 | assert out_path.exists() 48 | 49 | exported_model = torchvision_models.resnet18() 50 | exported_model.load_state_dict( 51 | torch.load(out_path, weights_only=True), strict=True 52 | ) 53 | -------------------------------------------------------------------------------- /tests/_models/torchvision/test_convnext.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import torch 9 | from torchvision import models 10 | 11 | from lightly_train._models.torchvision.convnext import ConvNeXtModelWrapper 12 | 13 | 14 | class TestConvNeXtModelWrapper: 15 | def test_feature_dim(self) -> None: 16 | model = models.convnext_tiny() 17 | feature_extractor = ConvNeXtModelWrapper(model=model) 18 | assert feature_extractor.feature_dim() == 768 19 | 20 | def test_forward_features(self) -> None: 21 | model = models.convnext_tiny() 22 | feature_extractor = ConvNeXtModelWrapper(model=model) 23 | x = torch.rand(1, 3, 224, 224) 24 | features = feature_extractor.forward_features(x)["features"] 25 | assert features.shape == (1, 768, 7, 7) 26 | 27 | def test_forward_pool(self) -> None: 28 | model = models.convnext_tiny() 29 | feature_extractor = ConvNeXtModelWrapper(model=model) 30 | x = torch.rand(1, 768, 7, 7) 31 | pool = feature_extractor.forward_pool({"features": x})["pooled_features"] 32 | assert pool.shape == (1, 768, 1, 1) 33 | 34 | def test_get_model(self) -> None: 35 | model = models.convnext_tiny() 36 | feature_extractor = ConvNeXtModelWrapper(model=model) 37 | assert feature_extractor.get_model() is model 38 | 39 | def test__device(self) -> None: 40 | # If this test fails it means the wrapped model doesn't move all required 41 | # modules to the correct device. This happens if not all required modules 42 | # are registered as attributes of the class. 43 | model = models.convnext_tiny() 44 | wrapped_model = ConvNeXtModelWrapper(model=model) 45 | wrapped_model.to("meta") 46 | wrapped_model.forward_features(torch.rand(1, 3, 224, 224, device="meta")) 47 | -------------------------------------------------------------------------------- /src/lightly_train/_transforms/predict_semantic_segmentation_transform.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from albumentations import ( 11 | BasicTransform, 12 | Compose, 13 | ) 14 | from albumentations.pytorch import ToTensorV2 15 | 16 | from lightly_train._transforms.predict_transform import ( 17 | PredictTransform, 18 | PredictTransformArgs, 19 | ) 20 | from lightly_train._transforms.transform import NormalizeArgs 21 | from lightly_train.types import TransformInput, TransformOutput 22 | 23 | 24 | class PredictSemanticSegmentationTransformArgs(PredictTransformArgs): 25 | image_size: tuple[int, int] 26 | normalize: NormalizeArgs 27 | 28 | 29 | class PredictSemanticSegmentationTransform(PredictTransform): 30 | transform_args_cls: type[PredictSemanticSegmentationTransformArgs] = ( 31 | PredictSemanticSegmentationTransformArgs 32 | ) 33 | 34 | def __init__( 35 | self, 36 | transform_args: PredictSemanticSegmentationTransformArgs, 37 | ) -> None: 38 | super().__init__(transform_args=transform_args) 39 | 40 | transform: list[BasicTransform] = [ 41 | # TODO(Yutong, 10/25): enable them once predict_batch is implemented 42 | # ToFloat(), 43 | # Normalize( 44 | # mean=transform_args.normalize.mean, 45 | # std=transform_args.normalize.std, 46 | # max_pixel_value=1.0, 47 | # ), 48 | # Resize( 49 | # height=transform_args.image_size[0], 50 | # width=transform_args.image_size[1], 51 | # ), 52 | ToTensorV2(), 53 | ] 54 | self.transform = Compose(transform) 55 | 56 | def __call__(self, input: TransformInput) -> TransformOutput: 57 | return [self.transform(**input)] 58 | -------------------------------------------------------------------------------- /src/lightly_train/_scaling.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from dataclasses import dataclass 11 | from typing import Iterable, TypeVar 12 | 13 | IMAGENET_SIZE = 1_000_000 # Approximate size of the ImageNet1k train dataset 14 | 15 | 16 | @dataclass 17 | class ScalingInfo: 18 | """Information used to scale method and optimizer parameters.""" 19 | 20 | dataset_size: int 21 | epochs: int 22 | 23 | 24 | def interpolate( 25 | input: float, 26 | input_start: float, 27 | input_end: float, 28 | value_start: float, 29 | value_end: float, 30 | round_ndigits: int | None = None, 31 | ) -> float: 32 | """Selects a value between value_start and value_end based on the input and the 33 | input range.""" 34 | value = value_start + (value_end - value_start) * (input - input_start) / ( 35 | input_end - input_start 36 | ) 37 | # Clamp the value to the range. 38 | value = max(value, value_start) 39 | value = min(value, value_end) 40 | return round(value, round_ndigits) 41 | 42 | 43 | _InputType = TypeVar("_InputType", bound=float) 44 | _ValueType = TypeVar("_ValueType") 45 | 46 | 47 | def get_bucket_value( 48 | input: _InputType, buckets: Iterable[tuple[_InputType, _ValueType]] 49 | ) -> _ValueType: 50 | """Map an input to a value based on a set of buckets. 51 | 52 | Args: 53 | input: 54 | The input value that is mapped to one of the buckets. 55 | buckets: 56 | A list of (threshold, value) tuples. 57 | 58 | Returns: 59 | The value from the first bucket that has a threshold greater than the input. 60 | """ 61 | for threshold, value in buckets: 62 | if input < threshold: 63 | return value 64 | raise ValueError(f"Input {input} is larger than all bucket thresholds.") 65 | -------------------------------------------------------------------------------- /tests/_models/test_embedding_model.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | 9 | import torch 10 | from torch import nn 11 | 12 | from lightly_train._models.embedding_model import EmbeddingModel 13 | 14 | from ..helpers import DummyCustomModel 15 | 16 | 17 | class TestEmbeddingModel: 18 | def test___init___with_embed_dim(self) -> None: 19 | wrapped_model = DummyCustomModel() 20 | embed_dim = 64 21 | model = EmbeddingModel(wrapped_model=wrapped_model, embed_dim=embed_dim) 22 | assert isinstance(model.embed_head, nn.Conv2d) 23 | assert model.embed_dim == embed_dim 24 | 25 | def test___init___without_embed_dim(self) -> None: 26 | wrapped_model = DummyCustomModel() 27 | model = EmbeddingModel(wrapped_model=wrapped_model, embed_dim=None) 28 | assert isinstance(model.embed_head, nn.Identity) 29 | assert model.embed_dim == wrapped_model.feature_dim() 30 | 31 | def test_forward__with_pooling(self) -> None: 32 | wrapped_model = DummyCustomModel() 33 | model = EmbeddingModel(wrapped_model=wrapped_model, embed_dim=64) 34 | x = torch.rand(2, 3, 32, 32) 35 | output = model(x, pool=True) 36 | assert output.shape == (2, 64, 1, 1) 37 | 38 | def test_forward__without_pooling(self) -> None: 39 | wrapped_model = DummyCustomModel() 40 | model = EmbeddingModel(wrapped_model=wrapped_model, embed_dim=64) 41 | x = torch.rand(2, 3, 32, 32) 42 | output = model(x, pool=False) 43 | assert output.shape == (2, 64, 31, 31) 44 | 45 | def test_forward__identity_head(self) -> None: 46 | wrapped_model = DummyCustomModel() 47 | model = EmbeddingModel(wrapped_model=wrapped_model, embed_dim=None) 48 | x = torch.rand(2, 3, 32, 32) 49 | output = model(x, pool=True) 50 | assert output.shape == (2, 2, 1, 1) 51 | -------------------------------------------------------------------------------- /tests/_models/torchvision/test_shufflenet.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | import torch 9 | from torchvision import models 10 | 11 | from lightly_train._models.torchvision.shufflenet import ShuffleNetV2ModelWrapper 12 | 13 | 14 | class TestShuffleNetV2ModelWrapper: 15 | def test_feature_dim(self) -> None: 16 | model = models.shufflenet_v2_x0_5() 17 | wrapped_model = ShuffleNetV2ModelWrapper(model=model) 18 | assert wrapped_model.feature_dim() == 1024 19 | 20 | def test_forward_features(self) -> None: 21 | model = models.shufflenet_v2_x0_5() 22 | wrapped_model = ShuffleNetV2ModelWrapper(model=model) 23 | x = torch.rand(1, 3, 224, 224) 24 | features = wrapped_model.forward_features(x)["features"] 25 | assert features.shape == (1, 1024, 7, 7) 26 | 27 | def test_forward_pool(self) -> None: 28 | model = models.shufflenet_v2_x0_5() 29 | wrapped_model = ShuffleNetV2ModelWrapper(model=model) 30 | x = torch.rand(1, 1024, 7, 7) 31 | pool = wrapped_model.forward_pool({"features": x})["pooled_features"] 32 | assert pool.shape == (1, 1024, 1, 1) 33 | 34 | def test_get_model(self) -> None: 35 | model = models.shufflenet_v2_x0_5() 36 | wrapped_model = ShuffleNetV2ModelWrapper(model=model) 37 | assert wrapped_model.get_model() is model 38 | 39 | def test__device(self) -> None: 40 | # If this test fails it means the wrapped model doesn't move all required 41 | # modules to the correct device. This happens if not all required modules 42 | # are registered as attributes of the class. 43 | model = models.shufflenet_v2_x0_5() 44 | wrapped_model = ShuffleNetV2ModelWrapper(model=model) 45 | wrapped_model.to("meta") 46 | wrapped_model.forward_features(torch.rand(1, 3, 224, 224, device="meta")) 47 | -------------------------------------------------------------------------------- /src/lightly_train/_models/package.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from abc import ABC, abstractmethod 11 | from pathlib import Path 12 | from typing import Any 13 | 14 | from lightly_train._models.model_wrapper import ModelWrapper 15 | from lightly_train.types import PackageModel 16 | 17 | 18 | class BasePackage(ABC): 19 | name: str # The name of the package. 20 | 21 | @classmethod 22 | @abstractmethod 23 | def export_model( 24 | cls, 25 | model: PackageModel | ModelWrapper | Any, 26 | out: Path, 27 | log_example: bool = True, 28 | ) -> None: 29 | """Export the model in the package's format. The model can be either a 30 | ModelWrapper or one of the package's underlying models. 31 | """ 32 | ... 33 | 34 | @classmethod 35 | @abstractmethod 36 | def is_supported_model(cls, model: PackageModel | ModelWrapper | Any) -> bool: 37 | """Check if the model is either a ModelWrapper or one of the package's 38 | underlying models. 39 | """ 40 | ... 41 | 42 | 43 | class Package(BasePackage): 44 | @classmethod 45 | @abstractmethod 46 | def list_model_names(cls) -> list[str]: 47 | """List all supported models by this package.""" 48 | ... 49 | 50 | @classmethod 51 | @abstractmethod 52 | def get_model( 53 | cls, 54 | model_name: str, 55 | num_input_channels: int = 3, 56 | model_args: dict[str, Any] | None = None, 57 | load_weights: bool = True, 58 | ) -> PackageModel: 59 | """Get the underlying model of the package by its name.""" 60 | ... 61 | 62 | @classmethod 63 | @abstractmethod 64 | def get_model_wrapper(cls, model: PackageModel) -> ModelWrapper: 65 | """Wrap the underlying model with the ModelWrapper.""" 66 | ... 67 | -------------------------------------------------------------------------------- /docs/source/data/multi_channel.md: -------------------------------------------------------------------------------- 1 | (multi-channel)= 2 | 3 | # Single- and Multi-Channel Images 4 | 5 | In addition to standard RGB images, LightlyTrain supports single- and multi-channel input for pretraining, and fine-tuning. 6 | 7 | ```{note} 8 | Multi-channel input is not supported for direct distillation because the DINOv2/v3 teacher models expect 3-channel input. However, you could load n-channel images and then reduce them to 3-channels with the [`ChannelDrop`](#method-transform-args-channel-drop) augmentation. 9 | ``` 10 | 11 | Specify the number of image channels and normalization parameteres in the respective LightlyTrain training function. For example, to fine-tune a semantic segmentation model on 4-channel images: 12 | 13 | ```python 14 | import lightly_train 15 | 16 | lightly_train.train_semantic_segmentation( 17 | out="out/my_experiment", 18 | model="dinov2/vitl14-eomt", 19 | data={ 20 | ... # multi-channel image data (e.g. RGB-NIR) 21 | }, 22 | transform_args={ 23 | "num_channels": 4, # specify number of channels here 24 | "normalize": { 25 | "mean": [0, 0, 0, 0], 26 | "std": [1, 1, 1, 1], 27 | }, 28 | }, 29 | ) 30 | ``` 31 | 32 | ## Models 33 | 34 | The following models support multi-channel image input: 35 | 36 | | Library | Supported Models | Docs | 37 | |---------|------------------|------| 38 | | LightlyTrain | DINOv3 | | 39 | | LightlyTrain | DINOv2 | | 40 | | TIMM | All models | [🔗](#models-timm) | 41 | 42 | ## Transforms 43 | 44 | The following image transforms are disabled for images that do not have 3 channels: 45 | 46 | - `ColorJitter` 47 | - `RandomGrayscale` 48 | - `Solarize` 49 | 50 | If any other transform defaults are incompatible with your data, you can disable them by setting the corresponding transform argument to `None`. For example, to disable `GaussianBlur`: 51 | 52 | ```python 53 | transform_args={ 54 | "num_channels": 4, 55 | "gaussian_blur": None 56 | }, 57 | ``` 58 | 59 | See [Configure Transform Arguments](#method-transform-args) for details on customizing transforms. 60 | -------------------------------------------------------------------------------- /src/lightly_train/_models/dinov3/dinov3_convnext.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from torch import Tensor 11 | from torch.nn import AdaptiveAvgPool2d, Module 12 | 13 | from lightly_train._models.dinov3.dinov3_src.models.convnext import ConvNeXt 14 | from lightly_train._models.model_wrapper import ( 15 | ForwardFeaturesOutput, 16 | ForwardPoolOutput, 17 | ModelWrapper, 18 | ) 19 | 20 | 21 | class DINOv3VConvNeXtModelWrapper(Module, ModelWrapper): 22 | def __init__(self, model: ConvNeXt) -> None: 23 | super().__init__() 24 | self._model = model 25 | self._feature_dim = int(self._model.embed_dim) 26 | self._pool = AdaptiveAvgPool2d((1, 1)) 27 | 28 | def feature_dim(self) -> int: 29 | return self._feature_dim 30 | 31 | def forward_features( 32 | self, x: Tensor, masks: Tensor | None = None 33 | ) -> ForwardFeaturesOutput: 34 | rt = self._model(x, masks, is_training=True) # forcing to return all patches 35 | if rt["x_norm_patchtokens"].dim() == 3: 36 | x_norm_patchtokens = rt["x_norm_patchtokens"] 37 | b = x_norm_patchtokens.shape[0] 38 | d = x_norm_patchtokens.shape[2] 39 | h, w = rt["x_norm_patchtokens_hw"] 40 | 41 | features_reshaped = x_norm_patchtokens.permute(0, 2, 1).reshape(b, d, h, w) 42 | else: 43 | raise ValueError( 44 | f"Unexpected shape for x_norm_patchtokens: {rt['x_norm_patchtokens'].shape}" 45 | ) 46 | return {"features": features_reshaped, "cls_token": rt["x_norm_clstoken"]} 47 | 48 | def forward_pool(self, x: ForwardFeaturesOutput) -> ForwardPoolOutput: 49 | return {"pooled_features": self._pool(x["features"])} 50 | 51 | def get_model(self) -> ConvNeXt: 52 | return self._model 53 | 54 | def make_teacher(self) -> None: 55 | pass 56 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to LightlyTrain 2 | 3 | ## Development 4 | 5 | ### Setting up the Development Environment 6 | 7 | ``` 8 | git clone https://github.com/lightly-ai/lightly-train.git 9 | uv venv .venv 10 | source .venv/bin/activate 11 | make install-dev 12 | ``` 13 | 14 | Make sure the environment is activated before running the following commands. 15 | 16 | > [!WARNING]\ 17 | > Prepending commands with `uv run` might not work properly. Activate the environment directly instead. 18 | 19 | ### Running Checks and Tests 20 | 21 | Before committing code, make sure all tests and checks pass: 22 | 23 | ``` 24 | make format 25 | make static-checks 26 | ``` 27 | 28 | and if you want to run all the tests: 29 | 30 | ``` 31 | make test 32 | ``` 33 | 34 | ### Documentation 35 | 36 | Documentation is in the [docs](./docs) folder. To build the documentation, install 37 | dev dependencies with `make install-dev`, then move to the `docs` folder and run: 38 | 39 | ``` 40 | make docs 41 | ``` 42 | 43 | This builds the documentation in the `docs/build/` folder. 44 | 45 | To build the documentation for the stable version, checkout the branch with the 46 | stable version and run: 47 | 48 | ``` 49 | make docs-stable 50 | ``` 51 | 52 | This builds the documentaion in the `docs/build/stable` folder. 53 | 54 | Docs can be served locally with: 55 | 56 | ``` 57 | make serve 58 | ``` 59 | 60 | #### Writing Documentation 61 | 62 | The documentation source is in [docs/source](./docs/source). The documentation is 63 | written in Markdown (MyST flavor). For more information regarding formatting, see: 64 | 65 | - https://pradyunsg.me/furo/reference/ 66 | - https://myst-parser.readthedocs.io/en/latest/syntax/typography.html 67 | 68 | ### Contributor License Agreement (CLA) 69 | 70 | To contribute to this repository, you must sign a Contributor License Agreement (CLA). 71 | This is a one-time process done through GitHub when you open your first pull request. 72 | You will be prompted automatically. 73 | 74 | By signing the CLA, you agree that your contributions may be used under the terms of the project license. 75 | -------------------------------------------------------------------------------- /src/lightly_train/_transforms/random_order.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | import numpy as np 11 | from albumentations import BaseCompose, BasicTransform, SomeOf 12 | from lightning_utilities.core.imports import RequirementCache 13 | from numpy.typing import NDArray 14 | 15 | ALBUMENTATIONS_GEQ_1_4_21 = RequirementCache("albumentations>=1.4.21") 16 | 17 | if not ALBUMENTATIONS_GEQ_1_4_21: 18 | from albumentations import random_utils 19 | 20 | 21 | class RandomOrder(SomeOf): # type: ignore[misc] 22 | def __init__( 23 | self, 24 | transforms: list[BasicTransform | BaseCompose], 25 | n: int | None = None, 26 | replace: bool = False, 27 | p: float = 1.0, 28 | ): 29 | """Apply a random number of transformations from a list in random order. 30 | 31 | Args: 32 | transforms: List of transformations to choose from. 33 | n: How many transformations to sample. If None, len(transforms) is used. 34 | replace: Whether to sample transformations with replacement. 35 | p: Probability of applying the entire pipeline, not each individual transform. 36 | """ 37 | if n is None: 38 | n = len(transforms) 39 | super().__init__(transforms=transforms, n=n, replace=replace, p=p) 40 | 41 | def _get_idx(self) -> NDArray[np.int64]: 42 | if ALBUMENTATIONS_GEQ_1_4_21: 43 | return self.random_generator.choice( # type: ignore[no-any-return] 44 | len(self.transforms), 45 | size=self.n, 46 | replace=self.replace, 47 | ) 48 | else: 49 | return random_utils.choice( # type: ignore[no-any-return] 50 | len(self.transforms), 51 | size=self.n, 52 | replace=self.replace, 53 | p=self.transforms_ps, 54 | ) 55 | -------------------------------------------------------------------------------- /docs/build.py: -------------------------------------------------------------------------------- 1 | # This script creates the build/index.html and build/versions.html files. 2 | # build/index.html redirects to build/stable/index.html. 3 | # build/versions.html lists all available documentation versions. 4 | 5 | import textwrap 6 | from argparse import ArgumentParser 7 | from pathlib import Path 8 | 9 | 10 | def build_index_html(build_dir: Path) -> None: 11 | """Creates the main index.html file that redirects to the stable version.""" 12 | html = textwrap.dedent(""" 13 | 14 | 15 | 16 | 17 | Redirecting... 18 | 19 | 20 |

If you are not redirected, click here.

21 | 22 | 23 | """) 24 | 25 | with open(build_dir / "index.html", "w") as f: 26 | f.write(html) 27 | 28 | 29 | def build_versions_html(build_dir: Path) -> None: 30 | """Creates the versions.html file that lists all available versions.""" 31 | 32 | header = textwrap.dedent(""" 33 | 34 | 35 | 36 | LightlyTrain Documentation 37 | 38 | 39 |

LightlyTrain Documentation

40 |
    41 | """) 42 | footer = textwrap.dedent(""" 43 |
44 | 45 | 46 | """) 47 | 48 | html = header 49 | versions = sorted( 50 | [path for path in build_dir.iterdir() if path.is_dir()], reverse=True 51 | ) 52 | for version in versions: 53 | html += f'
  • {version.name}
    • \n' 54 | html += footer 55 | 56 | with open(build_dir / "versions.html", "w") as f: 57 | f.write(html) 58 | 59 | 60 | def main(build_dir: Path) -> None: 61 | build_index_html(build_dir=build_dir) 62 | build_versions_html(build_dir=build_dir) 63 | 64 | 65 | if __name__ == "__main__": 66 | parser = ArgumentParser() 67 | parser.add_argument("--build-dir", type=Path, required=True) 68 | args = parser.parse_args() 69 | 70 | main(build_dir=args.build_dir) 71 | -------------------------------------------------------------------------------- /docs/source/pretrain_distill/models/timm.md: -------------------------------------------------------------------------------- 1 | (models-timm)= 2 | 3 | # TIMM 4 | 5 | This page describes how to use TIMM models with LightlyTrain. 6 | 7 | ```{important} 8 | [TIMM](https://github.com/huggingface/pytorch-image-models) must be installed with 9 | `pip install "lightly-train[timm]"`. 10 | ``` 11 | 12 | ## Pretrain and Fine-tune a TIMM Model 13 | 14 | ### Pretrain 15 | 16 | Pretraining TIMM models with LightlyTrain is straightforward. Below we provide the 17 | minimum scripts for pretraining using `timm/resnet18` as an example: 18 | 19 | ````{tab} Python 20 | ```python 21 | import lightly_train 22 | 23 | if __name__ == "__main__": 24 | lightly_train.pretrain( 25 | out="out/my_experiment", # Output directory. 26 | data="my_data_dir", # Directory with images. 27 | model="timm/resnet18", # Pass the timm model. 28 | ) 29 | 30 | ``` 31 | 32 | Or alternatively, pass directly a TIMM model instance: 33 | 34 | ```python 35 | import timm 36 | 37 | import lightly_train 38 | 39 | if __name__ == "__main__": 40 | model = timm.create_model("resnet18") # Load the model. 41 | lightly_train.pretrain( 42 | out="out/my_experiment", # Output directory. 43 | data="my_data_dir", # Directory with images. 44 | model=model, # Pass the TIMM model. 45 | ) 46 | ```` 47 | 48 | ````{tab} Command Line 49 | ```bash 50 | lightly-train pretrain out="out/my_experiment" data="my_data_dir" model="timm/resnet18" 51 | ```` 52 | 53 | ### Fine-tune 54 | 55 | After pretraining, you can load the exported model for fine-tuning with TIMM: 56 | 57 | ```python 58 | import timm 59 | 60 | model = timm.create_model( 61 | model_name="resnet18", 62 | checkpoint_path="out/my_experiment/exported_models/exported_last.pt", 63 | ) 64 | ``` 65 | 66 | ## Supported Models 67 | 68 | All timm models are supported, see [timm docs](https://github.com/huggingface/pytorch-image-models?tab=readme-ov-file#models) for a full list. 69 | 70 | Examples: 71 | 72 | - `timm/resnet50` 73 | - `timm/convnext_base` 74 | - `timm/vit_base_patch16_224` 75 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov2_eomt_semantic_segmentation/scheduler.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | from __future__ import annotations 9 | 10 | from torch.optim.lr_scheduler import LRScheduler 11 | from torch.optim.optimizer import Optimizer 12 | 13 | 14 | class TwoStageWarmupPolySchedule(LRScheduler): 15 | def __init__( 16 | self, 17 | optimizer: Optimizer, 18 | num_backbone_params: int, 19 | warmup_steps: tuple[int, int], 20 | total_steps: int, 21 | poly_power: float, 22 | last_epoch: int = -1, 23 | ): 24 | self.num_backbone_params = num_backbone_params 25 | self.warmup_steps = warmup_steps 26 | self.total_steps = total_steps 27 | self.poly_power = poly_power 28 | super().__init__(optimizer, last_epoch) 29 | 30 | def get_lr(self) -> list[float]: 31 | step = self.last_epoch 32 | lrs = [] 33 | non_vit_warmup, vit_warmup = self.warmup_steps 34 | for i, base_lr in enumerate(self.base_lrs): 35 | if i >= self.num_backbone_params: 36 | if non_vit_warmup > 0 and step < non_vit_warmup: 37 | lr = base_lr * (step / non_vit_warmup) 38 | else: 39 | adjusted = max(0, step - non_vit_warmup) 40 | max_steps = max(1, self.total_steps - non_vit_warmup) 41 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 42 | else: 43 | if step < non_vit_warmup: 44 | lr = 0 45 | elif step < non_vit_warmup + vit_warmup: 46 | lr = base_lr * ((step - non_vit_warmup) / vit_warmup) 47 | else: 48 | adjusted = max(0, step - non_vit_warmup - vit_warmup) 49 | max_steps = max(1, self.total_steps - non_vit_warmup - vit_warmup) 50 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 51 | lrs.append(lr) 52 | return lrs 53 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_instance_segmentation/scheduler.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | from __future__ import annotations 9 | 10 | from torch.optim.lr_scheduler import LRScheduler 11 | from torch.optim.optimizer import Optimizer 12 | 13 | 14 | class TwoStageWarmupPolySchedule(LRScheduler): 15 | def __init__( 16 | self, 17 | optimizer: Optimizer, 18 | num_backbone_params: int, 19 | warmup_steps: tuple[int, int], 20 | total_steps: int, 21 | poly_power: float, 22 | last_epoch: int = -1, 23 | ): 24 | self.num_backbone_params = num_backbone_params 25 | self.warmup_steps = warmup_steps 26 | self.total_steps = total_steps 27 | self.poly_power = poly_power 28 | super().__init__(optimizer, last_epoch) 29 | 30 | def get_lr(self) -> list[float]: 31 | step = self.last_epoch 32 | lrs = [] 33 | non_vit_warmup, vit_warmup = self.warmup_steps 34 | for i, base_lr in enumerate(self.base_lrs): 35 | if i >= self.num_backbone_params: 36 | if non_vit_warmup > 0 and step < non_vit_warmup: 37 | lr = base_lr * (step / non_vit_warmup) 38 | else: 39 | adjusted = max(0, step - non_vit_warmup) 40 | max_steps = max(1, self.total_steps - non_vit_warmup) 41 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 42 | else: 43 | if step < non_vit_warmup: 44 | lr = 0 45 | elif step < non_vit_warmup + vit_warmup: 46 | lr = base_lr * ((step - non_vit_warmup) / vit_warmup) 47 | else: 48 | adjusted = max(0, step - non_vit_warmup - vit_warmup) 49 | max_steps = max(1, self.total_steps - non_vit_warmup - vit_warmup) 50 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 51 | lrs.append(lr) 52 | return lrs 53 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_panoptic_segmentation/scheduler.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | from __future__ import annotations 9 | 10 | from torch.optim.lr_scheduler import LRScheduler 11 | from torch.optim.optimizer import Optimizer 12 | 13 | 14 | class TwoStageWarmupPolySchedule(LRScheduler): 15 | def __init__( 16 | self, 17 | optimizer: Optimizer, 18 | num_backbone_params: int, 19 | warmup_steps: tuple[int, int], 20 | total_steps: int, 21 | poly_power: float, 22 | last_epoch: int = -1, 23 | ): 24 | self.num_backbone_params = num_backbone_params 25 | self.warmup_steps = warmup_steps 26 | self.total_steps = total_steps 27 | self.poly_power = poly_power 28 | super().__init__(optimizer, last_epoch) 29 | 30 | def get_lr(self) -> list[float]: 31 | step = self.last_epoch 32 | lrs = [] 33 | non_vit_warmup, vit_warmup = self.warmup_steps 34 | for i, base_lr in enumerate(self.base_lrs): 35 | if i >= self.num_backbone_params: 36 | if non_vit_warmup > 0 and step < non_vit_warmup: 37 | lr = base_lr * (step / non_vit_warmup) 38 | else: 39 | adjusted = max(0, step - non_vit_warmup) 40 | max_steps = max(1, self.total_steps - non_vit_warmup) 41 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 42 | else: 43 | if step < non_vit_warmup: 44 | lr = 0 45 | elif step < non_vit_warmup + vit_warmup: 46 | lr = base_lr * ((step - non_vit_warmup) / vit_warmup) 47 | else: 48 | adjusted = max(0, step - non_vit_warmup - vit_warmup) 49 | max_steps = max(1, self.total_steps - non_vit_warmup - vit_warmup) 50 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 51 | lrs.append(lr) 52 | return lrs 53 | -------------------------------------------------------------------------------- /src/lightly_train/_task_models/dinov3_eomt_semantic_segmentation/scheduler.py: -------------------------------------------------------------------------------- 1 | # 2 | # --------------------------------------------------------------- 3 | # © 2025 Mobile Perception Systems Lab at TU/e. All rights reserved. 4 | # Licensed under the MIT License. 5 | # --------------------------------------------------------------- 6 | # 7 | 8 | from __future__ import annotations 9 | 10 | from torch.optim.lr_scheduler import LRScheduler 11 | from torch.optim.optimizer import Optimizer 12 | 13 | 14 | class TwoStageWarmupPolySchedule(LRScheduler): 15 | def __init__( 16 | self, 17 | optimizer: Optimizer, 18 | num_backbone_params: int, 19 | warmup_steps: tuple[int, int], 20 | total_steps: int, 21 | poly_power: float, 22 | last_epoch: int = -1, 23 | ): 24 | self.num_backbone_params = num_backbone_params 25 | self.warmup_steps = warmup_steps 26 | self.total_steps = total_steps 27 | self.poly_power = poly_power 28 | super().__init__(optimizer, last_epoch) 29 | 30 | def get_lr(self) -> list[float]: 31 | step = self.last_epoch 32 | lrs = [] 33 | non_vit_warmup, vit_warmup = self.warmup_steps 34 | for i, base_lr in enumerate(self.base_lrs): 35 | if i >= self.num_backbone_params: 36 | if non_vit_warmup > 0 and step < non_vit_warmup: 37 | lr = base_lr * (step / non_vit_warmup) 38 | else: 39 | adjusted = max(0, step - non_vit_warmup) 40 | max_steps = max(1, self.total_steps - non_vit_warmup) 41 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 42 | else: 43 | if step < non_vit_warmup: 44 | lr = 0 45 | elif step < non_vit_warmup + vit_warmup: 46 | lr = base_lr * ((step - non_vit_warmup) / vit_warmup) 47 | else: 48 | adjusted = max(0, step - non_vit_warmup - vit_warmup) 49 | max_steps = max(1, self.total_steps - non_vit_warmup - vit_warmup) 50 | lr = base_lr * (1 - (adjusted / max_steps)) ** self.poly_power 51 | lrs.append(lr) 52 | return lrs 53 | -------------------------------------------------------------------------------- /src/lightly_train/_callbacks/export.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | from __future__ import annotations 9 | 10 | from pathlib import Path 11 | 12 | from pytorch_lightning import LightningModule, Trainer 13 | from pytorch_lightning.callbacks import Callback 14 | from pytorch_lightning.utilities import rank_zero_only 15 | 16 | from lightly_train._commands import common_helpers 17 | from lightly_train._commands.common_helpers import ModelFormat 18 | from lightly_train._configs.config import PydanticConfig 19 | from lightly_train._models import package_helpers 20 | from lightly_train._models.model_wrapper import ModelWrapper 21 | 22 | 23 | class ModelExportArgs(PydanticConfig): 24 | every_n_epochs: int = 1 25 | 26 | 27 | class ModelExport(Callback): 28 | def __init__( 29 | self, 30 | wrapped_model: ModelWrapper, 31 | out_dir: Path, 32 | every_n_epochs: int = 1, 33 | ): 34 | self._wrapped_model = wrapped_model 35 | self._out_dir = out_dir 36 | self._every_n_epochs = every_n_epochs 37 | self._package = package_helpers.get_package_from_model( 38 | self._wrapped_model, include_custom=True, fallback_custom=True 39 | ) 40 | 41 | @rank_zero_only # type: ignore[misc] 42 | def _safe_export_model(self, export_path: Path) -> None: 43 | """Export the model to the specified path, deleting any existing file.""" 44 | if export_path.exists(): 45 | export_path.unlink(missing_ok=True) 46 | 47 | common_helpers.export_model( 48 | model=self._wrapped_model, 49 | out=export_path, 50 | format=ModelFormat.PACKAGE_DEFAULT, 51 | package=self._package, 52 | log_example=False, 53 | ) 54 | 55 | def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None: 56 | if pl_module.current_epoch % self._every_n_epochs == 0: 57 | # Delete the previous export if it exists 58 | export_path = self._out_dir / "exported_last.pt" 59 | self._safe_export_model(export_path) 60 | -------------------------------------------------------------------------------- /docker/Dockerfile-amd64-cuda: -------------------------------------------------------------------------------- 1 | # Image is released at https://hub.docker.com/r/lightly/train 2 | # Image is CUDA-optimized for single/multi-GPU training and inference 3 | 4 | # Start FROM PyTorch image https://hub.docker.com/r/pytorch/pytorch 5 | # We use CUDA 11.8 because it is compatible with older CUDA Drivers (>=450.80.02 linux, >=452.39 windows). 6 | # See https://docs.nvidia.com/deploy/cuda-compatibility/#minor-version-compatibility 7 | # The CUDA Driver is the only component from the host system that has to be compatible with the docker image. 8 | # The PyTorch and cuDNN versions are independent of the host system. 9 | # 10 | # The PyTorch 2.5.1 image comes with Python 3.11. 11 | FROM pytorch/pytorch:2.5.1-cuda11.8-cudnn9-runtime AS runtime 12 | 13 | # Install packages into the system Python and skip creating a virtual environment. 14 | ENV UV_SYSTEM_PYTHON="true" \ 15 | # Do not cache dependencies as they would also be saved in the docker image. 16 | UV_NO_CACHE="true" 17 | 18 | # Required for uv installation. 19 | RUN apt-get update && apt-get install -y make curl 20 | 21 | # Install Pillow-SIMD dependencies. 22 | RUN apt-get update && apt-get install -y python3.11-dev libjpeg8-dev libjpeg-turbo-progs libtiff5-dev libwebp-dev gcc 23 | 24 | # Create working directory 25 | WORKDIR /home/lightly_train 26 | 27 | # Set and create the directory to save pretrained torch models into 28 | ENV TORCH_HOME="/home/lightly_train/.cache/torch" 29 | RUN mkdir -p ${TORCH_HOME} && chmod -R a+w $TORCH_HOME 30 | 31 | # Set and create the directory to save pretrained models into 32 | ENV LIGHTLY_TRAIN_CACHE_DIR="/home/lightly_train/.cache/lightly-train" 33 | RUN mkdir -p ${LIGHTLY_TRAIN_CACHE_DIR} && chmod -R a+w $LIGHTLY_TRAIN_CACHE_DIR 34 | RUN WEIGHTS_PATH="${LIGHTLY_TRAIN_CACHE_DIR}/weights" && mkdir -p ${WEIGHTS_PATH} && chmod -R a+w $WEIGHTS_PATH 35 | 36 | # Install uv 37 | COPY Makefile /home/lightly_train 38 | RUN make install-uv 39 | 40 | # Add uv to PATH 41 | ENV PATH="/root/.local/bin:$PATH" 42 | 43 | # Install the package dependencies. 44 | COPY pyproject.toml /home/lightly_train 45 | RUN make install-docker-dependencies && make download-docker-models 46 | 47 | # Copy the package itself 48 | COPY src /home/lightly_train/src 49 | 50 | # Install the package. 51 | RUN make install-docker 52 | -------------------------------------------------------------------------------- /src/lightly_train/_callbacks/mlflow_logging.py: -------------------------------------------------------------------------------- 1 | # 2 | # Copyright (c) Lightly AG and affiliates. 3 | # All rights reserved. 4 | # 5 | # This source code is licensed under the license found in the 6 | # LICENSE file in the root directory of this source tree. 7 | # 8 | try: 9 | from mlflow.system_metrics.system_metrics_monitor import SystemMetricsMonitor 10 | except ImportError: 11 | SystemMetricsMonitor = None # type: ignore[misc, assignment] 12 | from pytorch_lightning import LightningModule, Trainer 13 | from pytorch_lightning.callbacks import Callback 14 | 15 | from lightly_train._configs.config import PydanticConfig 16 | from lightly_train._loggers.mlflow import MLFlowLogger 17 | 18 | 19 | class MLFlowLoggingArgs(PydanticConfig): 20 | pass 21 | 22 | 23 | class MLFlowLogging(Callback): 24 | def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None: 25 | self.system_monitor = None 26 | for logger in trainer.loggers: 27 | if isinstance(logger, MLFlowLogger) and SystemMetricsMonitor is not None: 28 | self.system_monitor = SystemMetricsMonitor( # type: ignore[no-untyped-call] 29 | run_id=logger.run_id, 30 | ) 31 | self.system_monitor.start() # type: ignore[no-untyped-call] 32 | with open(trainer.default_root_dir + "/train.log", "r") as _f: 33 | logger.experiment.log_text( 34 | run_id=logger.run_id, 35 | text=_f.read(), 36 | artifact_file="logs/train-start.log", 37 | ) 38 | break 39 | 40 | def on_fit_end(self, trainer: Trainer, pl_module: LightningModule) -> None: 41 | if self.system_monitor is not None: 42 | self.system_monitor.finish() # type: ignore[no-untyped-call] 43 | for logger in trainer.loggers: 44 | if isinstance(logger, MLFlowLogger): 45 | with open(trainer.default_root_dir + "/train.log", "r") as _f: 46 | logger.experiment.log_text( 47 | run_id=logger.run_id, 48 | text=_f.read(), 49 | artifact_file="logs/train-end.log", 50 | ) 51 | break 52 | --------------------------------------------------------------------------------