Better than PL auto strategy selection (#3)

* Better than PL auto strategy

* Update README

Author: hrukalive

README.md

@@ -5,15 +5,6 @@ A collection of neural vocoders suitable for singing voice synthesis tasks.
 ## If you have any questions, please open an issue.
 
 
-
-# 使用ddp
-```
-pl_trainer_strategy: 
-  name: ddp
-  process_group_backend: nccl
-  find_unused_parameters: true
-```
-
 # 预处理 
 python [process.py](process.py) --config 配置文件 --num_cpu 并行数量 --strx 1 代表 强制绝对路径 0 代表相对路径
 

train.py

@@ -61,7 +61,13 @@ def train(config, exp_name, work_dir):
         accelerator=config['pl_trainer_accelerator'],
         devices=config['pl_trainer_devices'],
         num_nodes=config['pl_trainer_num_nodes'],
-        strategy=get_strategy(config['pl_trainer_strategy']),
+        strategy=get_strategy(
+            config['pl_trainer_devices'],
+            config['pl_trainer_num_nodes'],
+            config['pl_trainer_accelerator'],
+            config['pl_trainer_strategy'],
+            config['pl_trainer_precision'],
+        ),
         precision=config['pl_trainer_precision'],
         callbacks=[
             DsModelCheckpoint(

utils/training_utils.py

@@ -304,16 +304,82 @@ def get_metrics(self, trainer, model):
         return items
 
 
-def get_strategy(strategy):
-    if strategy['name'] == 'auto':
-        return 'auto'
-
-    from lightning.pytorch.strategies import StrategyRegistry
-    if strategy['name'] not in StrategyRegistry:
-        available_names = ", ".join(sorted(StrategyRegistry.keys())) or "none"
-        raise ValueError(f"Invalid strategy name {strategy['name']}. Available names: {available_names}")
-
-    data = StrategyRegistry[strategy['name']]
-    params = data['init_params']
-    params.update({k: v for k, v in strategy.items() if k != 'name'})
-    return data['strategy'](**utils.filter_kwargs(params, data['strategy']))
+def get_strategy(
+    devices="auto",
+    num_nodes=1,
+    accelerator="auto",
+    strategy={"name": "auto"},
+    precision=None,
+):
+    from lightning.fabric.utilities.device_parser import _determine_root_gpu_device
+    from lightning.pytorch.accelerators import AcceleratorRegistry
+    from lightning.pytorch.accelerators.cuda import CUDAAccelerator
+    from lightning.pytorch.accelerators.mps import MPSAccelerator
+    from lightning.pytorch.strategies import Strategy, SingleDeviceStrategy, StrategyRegistry
+    from lightning.pytorch.trainer.connectors import accelerator_connector
+    from lightning.pytorch.utilities.rank_zero import rank_zero_warn
+    class _DsAcceleratorConnector(accelerator_connector._AcceleratorConnector):
+        def __init__(self) -> None:
+            accelerator_connector._register_external_accelerators_and_strategies()
+            self._registered_strategies = StrategyRegistry.available_strategies()
+            self._accelerator_types = AcceleratorRegistry.available_accelerators()
+            self._parallel_devices = []
+            self._check_config_and_set_final_flags(
+                strategy=strategy["name"],
+                accelerator=accelerator,
+                precision=precision,
+                plugins=[],
+                sync_batchnorm=False,
+            )
+            if self._accelerator_flag == "auto":
+                self._accelerator_flag = self._choose_auto_accelerator()
+            elif self._accelerator_flag == "gpu":
+                self._accelerator_flag = self._choose_gpu_accelerator_backend()
+            self._check_device_config_and_set_final_flags(devices=devices, num_nodes=num_nodes)
+            self._set_parallel_devices_and_init_accelerator()
+            if self._strategy_flag == "auto":
+                self._strategy_flag = self._choose_strategy()
+            self._check_strategy_and_fallback()
+            self._init_strategy()
+            for k in ["colossalai", "bagua", "hpu", "hpu_parallel", "hpu_single", "ipu", "ipu_strategy"]:
+                if k in StrategyRegistry:
+                    StrategyRegistry.remove(k)
+
+        def _init_strategy(self) -> None:
+            assert isinstance(self._strategy_flag, (str, Strategy))
+            if isinstance(self._strategy_flag, str):
+                if self._strategy_flag not in StrategyRegistry:
+                    available_names = ", ".join(sorted(StrategyRegistry.available_strategies())) or "none"
+                    raise KeyError(f"Invalid strategy name {strategy['name']}. Available names: {available_names}")
+                data = StrategyRegistry[self._strategy_flag]
+                params = {}
+                # Replicate additional logic for _choose_strategy when dealing with single device strategies
+                if issubclass(data["strategy"], SingleDeviceStrategy):
+                    if self._accelerator_flag == "hpu":
+                        params = {"device": torch.device("hpu")}
+                    elif self._accelerator_flag == "tpu":
+                        params = {"device": self._parallel_devices[0]}
+                    elif data["strategy"] is SingleDeviceStrategy:
+                        if isinstance(self._accelerator_flag, (CUDAAccelerator, MPSAccelerator)) or (
+                            isinstance(self._accelerator_flag, str) and self._accelerator_flag in ("cuda", "gpu", "mps")
+                        ):
+                            params = {"device": _determine_root_gpu_device(self._parallel_devices)}
+                        else:
+                            params = {"device": "cpu"}
+                    else:
+                        raise NotImplementedError
+                params.update(data["init_params"])
+                params.update({k: v for k, v in strategy.items() if k != "name"})
+                self.strategy = data["strategy"](**utils.filter_kwargs(params, data["strategy"]))
+            elif isinstance(self._strategy_flag, SingleDeviceStrategy):
+                params = {"device": self._strategy_flag.root_device}
+                params.update({k: v for k, v in strategy.items() if k != "name"})
+                self.strategy = self._strategy_flag.__class__(**utils.filter_kwargs(params, self._strategy_flag.__class__))
+            else:
+                rank_zero_warn(
+                    f"Inferred strategy {self._strategy_flag.__class__.__name__} cannot take custom configurations."
+                    f"To use custom configurations, please specify the strategy name explicitly."
+                )
+                self.strategy = self._strategy_flag
+
+    return _DsAcceleratorConnector().strategy