leave only parameters display

Author: OutisLi

deepmd/pt/train/training.py

@@ -743,63 +743,27 @@ def warm_up_linear(step: int, warmup_steps: int) -> float:
         self.profiling = training_params.get("profiling", False)
         self.profiling_file = training_params.get("profiling_file", "timeline.json")
 
-        # Log model summary info (descriptor type and parameter count)
+        # Log model parameter count
         if self.rank == 0:
-            self._log_model_summary()
-
-    def _log_model_summary(self) -> None:
-        """Log model summary information including descriptor type and parameter count."""
-
-        def get_descriptor_type(model: torch.nn.Module) -> str:
-            """Get the descriptor type name from model."""
-            # Standard models have get_descriptor method
-            if hasattr(model, "get_descriptor"):
-                descriptor = model.get_descriptor()
-                if descriptor is not None and hasattr(descriptor, "serialize"):
-                    serialized = descriptor.serialize()
-                    if isinstance(serialized, dict) and "type" in serialized:
-                        return str(serialized["type"]).upper()
-            # ZBL and other models: use serialize() API
-            if hasattr(model, "serialize"):
-                serialized = model.serialize()
-                if isinstance(serialized, dict):
-                    model_type = str(serialized.get("type", "")).lower()
-                    if model_type == "zbl":
-                        # ZBL model: get descriptor type from the DP sub-model
-                        models_data = serialized.get("models", [])
-                        if models_data and isinstance(models_data[0], dict):
-                            descriptor_data = models_data[0].get("descriptor", {})
-                            if isinstance(descriptor_data, dict):
-                                desc_type = descriptor_data.get("type", "UNKNOWN")
-                                return f"{str(desc_type).upper()} (with ZBL)"
-                        return "UNKNOWN (with ZBL)"
-            return "UNKNOWN"
-
-        def count_parameters(model: torch.nn.Module) -> tuple[int, int]:
-            """Count the number of trainable and total parameters.
-
-            Returns
-            -------
-            tuple[int, int]
-                (trainable_count, total_count)
-            """
-            trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-            total = sum(p.numel() for p in model.parameters())
-            return trainable, total
+            self._log_parameter_count()
 
+    def _log_parameter_count(self) -> None:
+        """Log model parameter count."""
         if not self.multi_task:
-            desc_type = get_descriptor_type(self.model)
-            trainable, total = count_parameters(self.model)
-            log.info(f"Descriptor:    {desc_type}")
+            trainable = sum(
+                p.numel() for p in self.model.parameters() if p.requires_grad
+            )
+            total = sum(p.numel() for p in self.model.parameters())
             log.info(
                 f"Model Params:  {total / 1e6:.3f} M   (Trainable: {trainable / 1e6:.3f} M)"
             )
         else:
-            # For multi-task, log each model's info
             for model_key in self.model_keys:
-                desc_type = get_descriptor_type(self.model[model_key])
-                trainable, total = count_parameters(self.model[model_key])
-                log.info(f"Descriptor [{model_key}]:   {desc_type}")
+                model = self.model[model_key]
+                trainable = sum(
+                    p.numel() for p in model.parameters() if p.requires_grad
+                )
+                total = sum(p.numel() for p in model.parameters())
                 log.info(
                     f"Model Params [{model_key}]: {total / 1e6:.3f} M   (Trainable: {trainable / 1e6:.3f} M)"
                 )