Spaces:
Configuration error
Configuration error
| 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 | |
| ) |