change BaseFinetuning.__apply_mapping_to_param_groups to protected instead of private and minor refactor of finetuning_callback test

Author: speediedan

pytorch_lightning/callbacks/finetuning.py

@@ -98,7 +98,7 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
         if self._restarting:
             named_parameters = dict(pl_module.named_parameters())
             for opt_idx, optimizer in enumerate(trainer.optimizers):
-                param_groups = self.__apply_mapping_to_param_groups(
+                param_groups = self._apply_mapping_to_param_groups(
                     self._internal_optimizer_metadata[opt_idx], named_parameters
                 )
                 optimizer.param_groups = param_groups
@@ -244,7 +244,7 @@ def on_before_accelerator_backend_setup(self, trainer: "pl.Trainer", pl_module:
         self.freeze_before_training(pl_module)
 
     @staticmethod
-    def __apply_mapping_to_param_groups(param_groups: List[Dict[str, Any]], mapping: dict) -> List[Dict[str, Any]]:
+    def _apply_mapping_to_param_groups(param_groups: List[Dict[str, Any]], mapping: dict) -> List[Dict[str, Any]]:
         output = []
         for g in param_groups:
             # skip params to save memory
@@ -262,13 +262,13 @@ def _store(
     ) -> None:
         mapping = {p: n for n, p in pl_module.named_parameters()}
         if opt_idx not in self._internal_optimizer_metadata:
-            self._internal_optimizer_metadata[opt_idx] = self.__apply_mapping_to_param_groups(
+            self._internal_optimizer_metadata[opt_idx] = self._apply_mapping_to_param_groups(
                 current_param_groups, mapping
             )
         elif num_param_groups != len(current_param_groups):
             # save new param_groups possibly created by the users.
             self._internal_optimizer_metadata[opt_idx].extend(
-                self.__apply_mapping_to_param_groups(current_param_groups[num_param_groups:], mapping)
+                self._apply_mapping_to_param_groups(current_param_groups[num_param_groups:], mapping)
             )
 
     def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:

tests/callbacks/test_finetuning_callback.py

@@ -287,34 +287,36 @@ def configure_optimizers(self):
     trainer.fit(model)
 
 
-def test_complex_nested_model():
-    """Test flattening, freezing, and thawing of models which contain parent (non-leaf) modules with parameters
-    directly themselves rather than exclusively their submodules containing parameters."""
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, 3)
+        self.act = nn.ReLU()
+        self.bn = nn.BatchNorm2d(out_channels)
 
-    class ConvBlock(nn.Module):
-        def __init__(self, in_channels, out_channels):
-            super().__init__()
-            self.conv = nn.Conv2d(in_channels, out_channels, 3)
-            self.act = nn.ReLU()
-            self.bn = nn.BatchNorm2d(out_channels)
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.act(x)
+        return self.bn(x)
 
-        def forward(self, x):
-            x = self.conv(x)
-            x = self.act(x)
-            return self.bn(x)
 
-    class ConvBlockParam(nn.Module):
-        def __init__(self, in_channels, out_channels):
-            super().__init__()
-            self.module_dict = nn.ModuleDict({"conv": nn.Conv2d(in_channels, out_channels, 3), "act": nn.ReLU()})
-            # add trivial test parameter to convblock to validate parent (non-leaf) module parameter handling
-            self.parent_param = nn.Parameter(torch.zeros((1), dtype=torch.float))
-            self.bn = nn.BatchNorm2d(out_channels)
+class ConvBlockParam(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.module_dict = nn.ModuleDict({"conv": nn.Conv2d(in_channels, out_channels, 3), "act": nn.ReLU()})
+        # add trivial test parameter to convblock to validate parent (non-leaf) module parameter handling
+        self.parent_param = nn.Parameter(torch.zeros((1), dtype=torch.float))
+        self.bn = nn.BatchNorm2d(out_channels)
 
-        def forward(self, x):
-            x = self.module_dict["conv"](x)
-            x = self.module_dict["act"](x)
-            return self.bn(x)
+    def forward(self, x):
+        x = self.module_dict["conv"](x)
+        x = self.module_dict["act"](x)
+        return self.bn(x)
+
+
+def test_complex_nested_model():
+    """Test flattening, freezing, and thawing of models which contain parent (non-leaf) modules with parameters
+    directly themselves rather than exclusively their submodules containing parameters."""
 
     model = nn.Sequential(
         OrderedDict(