Merge branch 'main' into trainers/segmentation-pretrained-weights

torchgeo/trainers/segmentation.py

@@ -95,6 +95,20 @@ def config_task(self) -> None:
                     state_dict = get_weight(weights).get_state_dict(progress=True)
                 self.model.encoder = utils.load_state_dict(self.model, state_dict)
 
+        # Freeze backbone
+        if self.hyperparams.get("freeze_backbone", False) and self.hyperparams[
+            "model"
+        ] in ["unet", "deeplabv3+"]:
+            for param in self.model.encoder.parameters():
+                param.requires_grad = False
+
+        # Freeze decoder
+        if self.hyperparams.get("freeze_decoder", False) and self.hyperparams[
+            "model"
+        ] in ["unet", "deeplabv3+"]:
+            for param in self.model.decoder.parameters():
+                param.requires_grad = False
+
     def __init__(self, **kwargs: Any) -> None:
         """Initialize the LightningModule with a model and loss function.