"examples/offline_inference/distributed.py" did not exist on "c5711ef98519de25d1f51121f7848a13f2891fc1"
interface.py 1.22 KB
Newer Older
1
import enum
2
from typing import Optional, Tuple
3

4
5
import torch

6
7
8
9

class PlatformEnum(enum.Enum):
    CUDA = enum.auto()
    ROCM = enum.auto()
10
    TPU = enum.auto()
11
    CPU = enum.auto()
12
    UNSPECIFIED = enum.auto()
13
14
15
16
17
18
19
20
21
22
23


class Platform:
    _enum: PlatformEnum

    def is_cuda(self) -> bool:
        return self._enum == PlatformEnum.CUDA

    def is_rocm(self) -> bool:
        return self._enum == PlatformEnum.ROCM

24
25
26
    def is_tpu(self) -> bool:
        return self._enum == PlatformEnum.TPU

27
28
29
    def is_cpu(self) -> bool:
        return self._enum == PlatformEnum.CPU

30
    @staticmethod
31
32
    def get_device_capability(device_id: int = 0) -> Optional[Tuple[int, int]]:
        return None
33

34
35
36
37
    @staticmethod
    def get_device_name(device_id: int = 0) -> str:
        raise NotImplementedError

38
39
40
41
42
43
44
45
46
47
48
49
50
    @staticmethod
    def inference_mode():
        """A device-specific wrapper of `torch.inference_mode`.

        This wrapper is recommended because some hardware backends such as TPU
        do not support `torch.inference_mode`. In such a case, they will fall
        back to `torch.no_grad` by overriding this method.
        """
        return torch.inference_mode(mode=True)


class UnspecifiedPlatform(Platform):
    _enum = PlatformEnum.UNSPECIFIED