FIX #123 trial 1: set_layer_trainable

peptdeep/model/model_interface.py

@@ -244,6 +244,28 @@ def build(self,
         self._model_to_device()
         self._init_for_training()
 
+    def set_bert_trainable(self,
+        trainable=False,  
+        bert_layer_name="hidden_nn",
+        bert_layer_idxes=[1,2,3], 
+    ):
+        self.set_layer_trainable(
+            trainable=trainable,
+            layer_names=[
+                f"{bert_layer_name}.bert.layer.{layer}"
+                for layer in bert_layer_idxes
+            ]
+        )
+
+    def set_layer_trainable(self, 
+        trainable=False, 
+        layer_names=[],
+    ):
+        for layer in layer_names:
+            self.model.get_submodule(
+                layer
+            ).requires_grad_(trainable)
+
     def train_with_warmup(self,
         precursor_df: pd.DataFrame,
         *,

peptdeep/model/ms2.py

@@ -209,7 +209,6 @@ def forward(self,
         NCEs:torch.Tensor,
         instrument_indices,
     ):
-
         in_x = self.dropout(self.input_nn(
             aa_indices, mod_x
         ))