refactor(pt): refactor training code

Author: iProzd

deepmd/pt/entrypoints/main.py

@@ -54,6 +54,7 @@
 from deepmd.pt.train import (
     training,
 )
+from deepmd.pt.train.trainer import Trainer as NewTrainer
 from deepmd.pt.train.wrapper import (
     ModelWrapper,
 )
@@ -106,6 +107,7 @@ def get_trainer(
     init_frz_model: str | None = None,
     shared_links: dict[str, Any] | None = None,
     finetune_links: dict[str, Any] | None = None,
+    use_legacy: bool = False,
 ) -> training.Trainer:
     multi_task = "model_dict" in config.get("model", {})
 
@@ -200,19 +202,34 @@ def prepare_trainer_input_single(
                 seed=data_seed,
             )
 
-    trainer = training.Trainer(
-        config,
-        train_data,
-        stat_file_path=stat_file_path,
-        validation_data=validation_data,
-        init_model=init_model,
-        restart_model=restart_model,
-        finetune_model=finetune_model,
-        force_load=force_load,
-        shared_links=shared_links,
-        finetune_links=finetune_links,
-        init_frz_model=init_frz_model,
-    )
+    if use_legacy:
+        trainer = training.Trainer(
+            config,
+            train_data,
+            stat_file_path=stat_file_path,
+            validation_data=validation_data,
+            init_model=init_model,
+            restart_model=restart_model,
+            finetune_model=finetune_model,
+            force_load=force_load,
+            shared_links=shared_links,
+            finetune_links=finetune_links,
+            init_frz_model=init_frz_model,
+        )
+    else:
+        trainer = NewTrainer(
+            config,
+            train_data,
+            stat_file_path=stat_file_path,
+            validation_data=validation_data,
+            init_model=init_model,
+            restart_model=restart_model,
+            finetune_model=finetune_model,
+            force_load=force_load,
+            shared_links=shared_links,
+            finetune_links=finetune_links,
+            init_frz_model=init_frz_model,
+        )
     return trainer
 
 

deepmd/pt/train/__init__.py

@@ -1 +1,102 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+"""PyTorch training module with modular, extensible design.
+
+This module provides a clean, component-based training system:
+
+- TrainingConfig: Configuration management with validation
+- DataManager: Data loading and batch iteration
+- OptimizerFactory: Strategy pattern for optimizer creation
+- CheckpointManager: Model persistence and recovery
+- TrainingLoop: Specialized training step implementations
+- HookManager: Extensible callback system
+- TrainingLogger: Formatted output and file I/O
+- Trainer: Main orchestrator coordinating all components
+
+Example:
+    >>> from deepmd.pt.train import Trainer, TrainingConfig
+    >>>
+    >>> # Create trainer
+    >>> trainer = Trainer(
+    ...     config=config_dict,
+    ...     training_data=train_dataset,
+    ...     validation_data=valid_dataset,
+    ... )
+    >>>
+    >>> # Run training
+    >>> trainer.run()
+
+Future extensions for multi-backend support:
+- AbstractTrainingLoop can be extended for JAX/NumPy
+- OptimizerFactory can support backend-specific optimizers
+- DataManager can use backend-specific data loading
+"""
+
+from deepmd.pt.train.checkpoint_manager import (
+    CheckpointManager,
+)
+from deepmd.pt.train.config import (
+    CheckpointConfig,
+    DisplayConfig,
+    LearningRateConfig,
+    OptimizerConfig,
+    TrainingConfig,
+)
+from deepmd.pt.train.data_manager import (
+    DataManager,
+)
+from deepmd.pt.train.hooks import (
+    HookManager,
+    HookPriority,
+    TensorBoardHook,
+    TimingHook,
+    TrainingHook,
+)
+from deepmd.pt.train.logger import (
+    LossAccumulator,
+    TrainingLogger,
+)
+from deepmd.pt.train.optimizer_factory import (
+    OptimizerFactory,
+)
+from deepmd.pt.train.trainer import (
+    Trainer,
+)
+
+# Keep old Trainer available for backward compatibility during transition
+from deepmd.pt.train.training import Trainer as LegacyTrainer
+from deepmd.pt.train.training_loop import (
+    AdamTrainingLoop,
+    BaseTrainingLoop,
+    LKFEnergyTrainingLoop,
+    TrainingLoopFactory,
+)
+from deepmd.pt.train.wrapper import (
+    ModelWrapper,
+)
+
+__all__ = [
+    # New modular components
+    "AdamTrainingLoop",
+    "BaseTrainingLoop",
+    "CheckpointConfig",
+    "CheckpointManager",
+    "DataManager",
+    "DisplayConfig",
+    "HookManager",
+    "HookPriority",
+    "LKFEnergyTrainingLoop",
+    "LearningRateConfig",
+    # Legacy support
+    "LegacyTrainer",
+    "LossAccumulator",
+    "ModelWrapper",
+    "OptimizerConfig",
+    "OptimizerFactory",
+    "TensorBoardHook",
+    "TimingHook",
+    "Trainer",
+    "TrainingConfig",
+    "TrainingHook",
+    "TrainingLogger",
+    "TrainingLoopFactory",
+]