1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
from typing import List, Tuple, Type, TYPE_CHECKING, Union
from ignite.distributed.comp_models.base import _SerialModel
from ignite.distributed.comp_models.horovod import has_hvd_support
from ignite.distributed.comp_models.native import has_native_dist_support
from ignite.distributed.comp_models.xla import has_xla_support
if TYPE_CHECKING:
from ignite.distributed.comp_models.horovod import _HorovodDistModel
from ignite.distributed.comp_models.native import _NativeDistModel
from ignite.distributed.comp_models.xla import _XlaDistModel
def setup_available_computation_models() -> (
Tuple[Type[Union[_SerialModel, "_NativeDistModel", "_XlaDistModel", "_HorovodDistModel"]], ...]
):
models: List[Type[Union[_SerialModel, "_NativeDistModel", "_XlaDistModel", "_HorovodDistModel"]]] = [
_SerialModel,
]
if has_native_dist_support:
from ignite.distributed.comp_models.native import _NativeDistModel
models.append(_NativeDistModel)
if has_xla_support:
from ignite.distributed.comp_models.xla import _XlaDistModel
models.append(_XlaDistModel)
if has_hvd_support:
from ignite.distributed.comp_models.horovod import _HorovodDistModel
models.append(_HorovodDistModel)
return tuple(models)
registered_computation_models = setup_available_computation_models()
|
