Add FastViT

Makefile

@@ -154,8 +154,7 @@ add-header:
 
 	# Apply the Apache 2.0 license header to DEIMv2 derived files
 	uv run --frozen licenseheaders -t dev_tools/deimv2_licenseheader.tmpl \
-		-f src/lightly_train/_task_models/dinov2_ltdetr_object_detection/dinov2_vit_wrapper.py \
-		src/lightly_train/_task_models/dinov3_ltdetr_object_detection/dinov3_vit_wrapper.py \
+		-f src/lightly_train/_task_models/dinov3_ltdetr_object_detection/dinov3_vit_wrapper.py \
 		src/lightly_train/_task_models/object_detection_components/flat_cosine.py \
 		-E py
 

pyproject.toml

@@ -37,18 +37,19 @@ dependencies = [
     "nvidia-ml-py>=12", # For GPU monitoring
     "pytorch_lightning>=2.1",
     "tensorboard>=2.10.0",
-    "torch>=2.1.0",    
+    "torch>=2.1.0",
     "torchmetrics[detection]>=0.8",    # Torchmetrics is needed for fine-tuning.
-                            # 0.8 is required by SuperGradients
-                            # 1.5.2 is the last version that supports Python 3.8
-                            # 1.7 is recommended version
+    # 0.8 is required by SuperGradients
+    # 1.5.2 is the last version that supports Python 3.8
+    # 1.7 is recommended version
     "torchvision>=0.16.0",  # torchvision 0.15.0 was yanked
     "transformers>=4.46",   # Transformers is needed for fine-tuning.
-                            # 4.46 is last version that supports Python 3.8.
-                            # 4.51.3 is the default for EoMT
+    # 4.46 is last version that supports Python 3.8.
+    # 4.51.3 is the default for EoMT
     "tqdm>=4.0.0",
     "pycocotools>=2.0.11; python_version>='3.9'",
     "numpy>1.23.0; python_version>='3.12'",  # numpy<=1.23 fails to build on Python>3.12 (no distutils).
+
 ]
 
 [dependency-groups]
@@ -190,7 +191,8 @@ ignore_missing_imports = true
 module = [
     "lightly_train._models.dinov2_vit.dinov2_vit_src.*",
     "lightly_train._models.dinov3.dinov3_src.*",
-    "lightly_train._task_models.object_detection_components.*"
+    "lightly_train._task_models.object_detection_components.*",
+    "lightly_train._models.fastvit.components.*",
 ]
 ignore_errors = true
 

src/lightly_train/_commands/train_task_helpers.py

@@ -61,9 +61,6 @@
 from lightly_train._task_models.dinov2_linear_semantic_segmentation.train_model import (
     DINOv2LinearSemanticSegmentationTrain,
 )
-from lightly_train._task_models.dinov2_ltdetr_object_detection.train_model import (
-    DINOv2LTDETRObjectDetectionTrain,
-)
 from lightly_train._task_models.dinov3_eomt_instance_segmentation.train_model import (
     DINOv3EoMTInstanceSegmentationTrain,
 )
@@ -129,7 +126,6 @@
     DINOv2LinearSemanticSegmentationTrain,
     DINOv3EoMTSemanticSegmentationTrain,
     SemanticSegmentationMultiheadTrain,
-    DINOv2LTDETRObjectDetectionTrain,
     DINOv3LTDETRObjectDetectionTrain,
     PicoDetObjectDetectionTrain,
 ]

src/lightly_train/_models/dinov2_vit/dinov2_vit.py

@@ -22,13 +22,11 @@
     ArchitectureInfoGettable,
     ForwardFeaturesOutput,
     ForwardPoolOutput,
-    MultiScaleFeatureModelWrapper,
+    MultiScaleFeatureViT,
 )
 
 
-class DINOv2ViTModelWrapper(
-    Module, MultiScaleFeatureModelWrapper, ArchitectureInfoGettable
-):
+class DINOv2ViTModelWrapper(Module, MultiScaleFeatureViT, ArchitectureInfoGettable):
     def __init__(self, model: DinoVisionTransformer) -> None:
         super().__init__()
         self._model = model
@@ -37,6 +35,9 @@ def __init__(self, model: DinoVisionTransformer) -> None:
     def feature_dim(self) -> int:
         return self._feature_dim
 
+    def patch_size(self) -> int:
+        return int(self._model.patch_size)
+
     def forward_features(
         self, x: Tensor, masks: Tensor | None = None, n_blocks: int = 1
     ) -> ForwardFeaturesOutput:

src/lightly_train/_models/dinov3/dinov3_convnext.py

@@ -12,17 +12,17 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor
-from torch.nn import AdaptiveAvgPool2d, Module
+from torch.nn import AdaptiveAvgPool2d, Conv2d, Module
 
 from lightly_train._models.dinov3.dinov3_src.models.convnext import ConvNeXt
 from lightly_train._models.model_wrapper import (
     ForwardFeaturesOutput,
     ForwardPoolOutput,
-    MultiScaleFeatureModelWrapper,
+    MultiScaleFeatureCNN,
 )
 
 
-class DINOv3VConvNeXtModelWrapper(Module, MultiScaleFeatureModelWrapper):
+class DINOv3VConvNeXtModelWrapper(Module, MultiScaleFeatureCNN):
     def __init__(self, model: ConvNeXt) -> None:
         super().__init__()
         self._model = model
@@ -83,6 +83,16 @@ def get_model(self) -> ConvNeXt:
     def make_teacher(self) -> None:
         pass
 
+    def multiscale_feature_strides(self) -> list[int]:
+        strides = []
+        cumulative = 1
+        for stage in self._model.downsample_layers:
+            for mod in stage:
+                if isinstance(mod, Conv2d):
+                    cumulative *= mod.stride[0]
+            strides.append(cumulative)
+        return strides
+
     def multiscale_feature_dims(self) -> list[int]:
         return list(self._model.embed_dims)
 

src/lightly_train/_models/dinov3/dinov3_src/models/convnext.py

@@ -158,7 +158,7 @@ def __init__(
         del ignored_kwargs
 
         # ==== ConvNeXt's original init =====
-        self.downsample_layers = (
+        self.downsample_layers: Sequence[nn.Sequential] = (
             nn.ModuleList()
         )  # stem and 3 intermediate downsampling conv layers
         stem = nn.Sequential(

src/lightly_train/_models/dinov3/dinov3_vit.py

@@ -24,13 +24,11 @@
     ArchitectureInfoGettable,
     ForwardFeaturesOutput,
     ForwardPoolOutput,
-    MultiScaleFeatureModelWrapper,
+    MultiScaleFeatureViT,
 )
 
 
-class DINOv3ViTModelWrapper(
-    Module, MultiScaleFeatureModelWrapper, ArchitectureInfoGettable
-):
+class DINOv3ViTModelWrapper(Module, MultiScaleFeatureViT, ArchitectureInfoGettable):
     def __init__(self, model: DinoVisionTransformer) -> None:
         super().__init__()
         self._model = model
@@ -39,6 +37,9 @@ def __init__(self, model: DinoVisionTransformer) -> None:
     def feature_dim(self) -> int:
         return self._feature_dim
 
+    def patch_size(self) -> int:
+        return int(self._model.patch_size)
+
     def forward_features(
         self, x: Tensor, masks: Tensor | None = None, n_blocks: int = 1
     ) -> ForwardFeaturesOutput:

src/lightly_train/_models/fastvit/components/__init__.py

@@ -0,0 +1,7 @@
+#
+# Copyright (c) Lightly AG and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#

src/lightly_train/_models/fastvit/components/models/__init__.py

@@ -0,0 +1,26 @@
+#
+# Copyright (c) Lightly AG and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+from .fastvit import (
+    fastvit_ma36,
+    fastvit_s12,
+    fastvit_sa12,
+    fastvit_sa24,
+    fastvit_sa36,
+    fastvit_t8,
+    fastvit_t12,
+)
+
+__all__ = [
+    "fastvit_t8",
+    "fastvit_t12",
+    "fastvit_s12",
+    "fastvit_sa12",
+    "fastvit_sa24",
+    "fastvit_sa36",
+    "fastvit_ma36",
+]