Add trainable settings for pt

deepmd/pt/model/descriptor/dpa1.py

@@ -70,8 +70,6 @@
             raise NotImplementedError("type_one_side is not supported.")
         if precision != "default" and precision != "float64":
             raise NotImplementedError("precison is not supported.")
-        if not trainable:
-            raise NotImplementedError("trainable == False is not supported.")
         if exclude_types is not None and exclude_types != []:
             raise NotImplementedError("exclude_types is not supported.")
         if stripped_type_embedding:
@@ -106,6 +104,9 @@
         self.type_embedding = TypeEmbedNet(ntypes, tebd_dim)
         self.tebd_dim = tebd_dim
         self.concat_output_tebd = concat_output_tebd
+        # set trainable
+        for param in self.parameters():
+            param.requires_grad = trainable
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""

deepmd/pt/model/descriptor/dpa2.py

@@ -75,6 +75,7 @@
         repformer_update_style: str = "res_avg",
         repformer_set_davg_zero: bool = True,  # TODO
         repformer_add_type_ebd_to_seq: bool = False,
+        trainable: bool = True,
         type: Optional[
             str
         ] = None,  # work around the bad design in get_trainer and DpLoaderSet!
@@ -170,6 +171,8 @@
             repformers block: set the avg to zero in statistics
         repformer_add_type_ebd_to_seq : bool
             repformers block: concatenate the type embedding at the output.
+        trainable : bool
+            If the parameters in the descriptor are trainable.
 
         Returns
         -------
@@ -249,6 +252,9 @@
         self.rcut = self.repinit.get_rcut()
         self.ntypes = ntypes
         self.sel = self.repinit.sel
+        # set trainable
+        for param in self.parameters():
+            param.requires_grad = trainable
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""

deepmd/pt/model/descriptor/se_a.py

@@ -272,6 +272,7 @@
         exclude_types: List[Tuple[int, int]] = [],
         old_impl: bool = False,
         type_one_side: bool = True,
+        trainable: bool = True,
         **kwargs,
     ):
         """Construct an embedding net of type `se_a`.
@@ -344,6 +345,9 @@
                 )
             self.filter_layers = filter_layers
         self.stats = None
+        # set trainable
+        for param in self.parameters():
+            param.requires_grad = trainable
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""

deepmd/pt/model/descriptor/se_r.py

@@ -63,6 +63,7 @@
         resnet_dt: bool = False,
         exclude_types: List[Tuple[int, int]] = [],
         old_impl: bool = False,
+        trainable: bool = True,
         **kwargs,
     ):
         super().__init__()
@@ -110,6 +111,9 @@
             )
         self.filter_layers = filter_layers
         self.stats = None
+        # set trainable
+        for param in self.parameters():
+            param.requires_grad = trainable
 
     def get_rcut(self) -> float:
         """Returns the cut-off radius."""

deepmd/pt/model/task/fitting.py

@@ -246,6 +246,8 @@
         Random seed.
     exclude_types: List[int]
         Atomic contributions of the excluded atom types are set zero.
+    trainable : bool
+            If the parameters in the fitting net are trainable.
 
     """
 
@@ -265,6 +267,7 @@
         rcond: Optional[float] = None,
         seed: Optional[int] = None,
         exclude_types: List[int] = [],
+        trainable: bool = True,
         **kwargs,
     ):
         super().__init__()
@@ -282,6 +285,11 @@
         self.rcond = rcond
         # order matters, should be place after the assignment of ntypes
         self.reinit_exclude(exclude_types)
+        self.trainable = trainable
+        # need support for each layer settings
+        self.trainable = (
+            all(self.trainable) if isinstance(self.trainable, list) else self.trainable
+        )
 
         net_dim_out = self._net_out_dim()
         # init constants
@@ -356,6 +364,9 @@
         if seed is not None:
             log.info("Set seed to %d in fitting net.", seed)
             torch.manual_seed(seed)
+        # set trainable
+        for param in self.parameters():
+            param.requires_grad = self.trainable
 
     def reinit_exclude(
         self,
@@ -397,7 +408,7 @@
             # "spin": self.spin ,
             ## NOTICE:  not supported by far
             "tot_ener_zero": False,
-            "trainable": [True] * (len(self.neuron) + 1),
+            "trainable": [self.trainable] * (len(self.neuron) + 1),
             "layer_name": None,
             "use_aparam_as_mask": False,
             "spin": None,

deepmd/pt/train/training.py

@@ -399,9 +399,6 @@ def get_loss(loss_params, start_lr, _ntypes):
             frz_model = torch.jit.load(init_frz_model, map_location=DEVICE)
             self.model.load_state_dict(frz_model.state_dict())
 
-        # Set trainable params
-        self.wrapper.set_trainable_params()
-
         # Multi-task share params
         if shared_links is not None:
             self.wrapper.share_params(shared_links, resume=model_params["resuming"])

deepmd/pt/train/wrapper.py

@@ -60,28 +60,6 @@ def __init__(
                     self.loss[task_key] = loss[task_key]
         self.inference_only = self.loss is None
 
-    def set_trainable_params(self):
-        supported_types = ["type_embedding", "descriptor", "fitting_net"]
-        for model_item in self.model:
-            for net_type in supported_types:
-                trainable = True
-                if not self.multi_task:
-                    if net_type in self.model_params:
-                        trainable = self.model_params[net_type].get("trainable", True)
-                else:
-                    if net_type in self.model_params["model_dict"][model_item]:
-                        trainable = self.model_params["model_dict"][model_item][
-                            net_type
-                        ].get("trainable", True)
-                if (
-                    hasattr(self.model[model_item], net_type)
-                    and getattr(self.model[model_item], net_type) is not None
-                ):
-                    for param in (
-                        self.model[model_item].__getattr__(net_type).parameters()
-                    ):
-                        param.requires_grad = trainable
-
     def share_params(self, shared_links, resume=False):
         supported_types = ["type_embedding", "descriptor", "fitting_net"]
         for shared_item in shared_links: