pytorch-lightning

1
# Copyright The Lightning AI team.
2
#
3
# Licensed under the Apache License, Version 2.0 (the "License");
4
# you may not use this file except in compliance with the License.
5
# You may obtain a copy of the License at
6
#
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
9
# Unless required by applicable law or agreed to in writing, software
10
# distributed under the License is distributed on an "AS IS" BASIS,
11
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
# See the License for the specific language governing permissions and
13
# limitations under the License.
14
import sys
15

16
from lightning.fabric.strategies.registry import _StrategyRegistry
17
from lightning.fabric.utilities.registry import _register_classes
18
from lightning.pytorch.strategies.ddp import DDPStrategy
19
from lightning.pytorch.strategies.deepspeed import DeepSpeedStrategy
20
from lightning.pytorch.strategies.fsdp import FSDPStrategy
21
from lightning.pytorch.strategies.parallel import ParallelStrategy
22
from lightning.pytorch.strategies.single_device import SingleDeviceStrategy
23
from lightning.pytorch.strategies.single_xla import SingleDeviceXLAStrategy  # noqa: F401
24
from lightning.pytorch.strategies.strategy import Strategy
25
from lightning.pytorch.strategies.xla import XLAStrategy  # noqa: F401
26

27
StrategyRegistry = _StrategyRegistry()
28
_register_classes(StrategyRegistry, "register_strategies", sys.modules[__name__], Strategy)
29

30
__all__ = [
31
    "DDPStrategy",
32
    "DeepSpeedStrategy",
33
    "FSDPStrategy",
34
    "ParallelStrategy",
35
    "SingleDeviceStrategy",
36
    "Strategy",
37
]
38