add testing

Author: SkafteNicki

tests/tests_pytorch/tuner/test_lr_finder.py

@@ -25,6 +25,7 @@
 
 from lightning.pytorch import Trainer, seed_everything
 from lightning.pytorch.callbacks import EarlyStopping
+from lightning.pytorch.callbacks.finetuning import BackboneFinetuning
 from lightning.pytorch.callbacks.lr_finder import LearningRateFinder
 from lightning.pytorch.demos.boring_classes import BoringModel
 from lightning.pytorch.tuner.lr_finder import _LRFinder
@@ -800,3 +801,47 @@ def configure_optimizers(self):
     assert len(lr_find_checkpoints) == 0, (
         f"lr_find checkpoint files should be cleaned up, but found: {lr_find_checkpoints}"
     )
+
+
+def test_lr_finder_with_backbone_finetuning_callback(tmp_path):
+    """Test that lr_find works correctly with BackboneFinetuning callback."""
+
+    class ModelWithBackbone(BoringModel):
+        def __init__(self):
+            super().__init__()
+            # Create a simple backbone-head architecture
+            self.backbone = torch.nn.Sequential(torch.nn.Linear(32, 16), torch.nn.ReLU(), torch.nn.Linear(16, 8))
+            self.head = torch.nn.Linear(8, 2)
+            self.learning_rate = 1e-3
+
+        def forward(self, x):
+            backbone_features = self.backbone(x)
+            return self.head(backbone_features)
+
+        def configure_optimizers(self):
+            return torch.optim.Adam(filter(lambda p: p.requires_grad, self.parameters()), lr=self.learning_rate)
+
+    model = ModelWithBackbone()
+    backbone_finetuning = BackboneFinetuning(unfreeze_backbone_at_epoch=1)
+
+    trainer = Trainer(
+        default_root_dir=tmp_path,
+        max_epochs=3,
+        enable_checkpointing=False,
+        enable_progress_bar=False,
+        enable_model_summary=False,
+        logger=False,
+        callbacks=[backbone_finetuning],
+    )
+
+    tuner = Tuner(trainer)
+    lr_finder = tuner.lr_find(model, num_training=5)
+
+    assert lr_finder is not None
+    assert hasattr(lr_finder, "results")
+    assert len(lr_finder.results) > 0
+    trainer.fit(model)
+
+    # Check that backbone was unfrozen at the correct epoch
+    for param in model.backbone.parameters():
+        assert param.requires_grad, "Backbone parameters should be unfrozen after epoch 1"