| import torch |
| from multiprocessing import cpu_count |
|
|
| class Config: |
| """Конфигурация для VC""" |
| |
| def __init__(self, device_str: str = "cuda" if torch.cuda.is_available() else "cpu") -> None: |
| """ |
| Инициализация конфигурации |
| |
| Args: |
| device_str: Строка устройства |
| """ |
| self.device_str: str = device_str |
| self.device_ids = None |
| self.set_device(self.device_str) |
| self.is_half: bool = False |
| self.n_cpu: int = cpu_count() |
| self.gpu_name = None |
| self.gpu_mem = None |
| self.x_pad, self.x_query, self.x_center, self.x_max = self.device_config() |
|
|
| def set_device(self, device_str: str) -> None: |
| """ |
| Установить устройство |
| |
| Args: |
| device_str: Строка устройства |
| """ |
| if "cuda" in device_str.lower(): |
| if ":" in device_str: |
| device_spec = device_str.split(":")[1] |
| self.device_ids = [int(id) for id in device_spec.split(",") if id.isdigit()] |
| else: |
| self.device_ids = list(range(torch.cuda.device_count())) |
| self.device = torch.device("cuda" if not self.device_ids else f"cuda:{self.device_ids[0]}") |
| elif "mps" in device_str.lower(): |
| self.device_ids = None |
| self.device = torch.device("mps") |
| else: |
| self.device_ids = None |
| self.device = torch.device("cpu") |
|
|
| def device_config(self): |
| """ |
| Настройка параметров для устройства |
| |
| Returns: |
| Кортеж (x_pad, x_query, x_center, x_max) |
| """ |
| if self.device.type == "cuda": |
| if self.device_ids: |
| self.gpu_mem = self._configure_gpu(self.device_ids[0]) |
|
|
| x_pad, x_query, x_center, x_max = ( |
| (3, 10, 60, 65) if self.is_half else (1, 6, 38, 41) |
| ) |
| if self.gpu_mem is not None and self.gpu_mem <= 4: |
| x_pad, x_query, x_center, x_max = (1, 5, 30, 32) |
|
|
| return x_pad, x_query, x_center, x_max |
|
|
| def _configure_gpu(self, device_id: int) -> int: |
| """ |
| Настройка GPU |
| |
| Args: |
| device_id: ID устройства |
| |
| Returns: |
| Объем памяти GPU в GB |
| """ |
| self.gpu_name = torch.cuda.get_device_name(f"cuda:{device_id}") |
| low_end_gpus = ["16", "P40", "P10", "1060", "1070", "1080"] |
| if ( |
| any(gpu in self.gpu_name for gpu in low_end_gpus) |
| and "V100" not in self.gpu_name.upper() |
| ): |
| self.is_half = False |
| return int( |
| torch.cuda.get_device_properties(self.device).total_memory |
| / 1024 |
| / 1024 |
| / 1024 |
| + 0.4 |
| ) |
