mirror of
https://github.com/iperov/DeepFaceLive
synced 2025-07-08 05:51:41 -07:00
66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
from typing import List
|
|
|
|
import cupy as cp
|
|
|
|
|
|
class CuPyDeviceInfo:
|
|
"""
|
|
Represents picklable CuPy device info
|
|
"""
|
|
def __init__(self, index=None, name=None, total_memory=None):
|
|
self._index : int = index
|
|
self._name : str = name
|
|
self._total_memory : int = total_memory
|
|
|
|
def __getstate__(self):
|
|
return self.__dict__.copy()
|
|
|
|
def __setstate__(self, d):
|
|
self.__init__()
|
|
self.__dict__.update(d)
|
|
|
|
def is_cpu(self) -> bool: return self._index == -1
|
|
|
|
def get_index(self) -> int:
|
|
return self._index
|
|
|
|
def get_name(self) -> str:
|
|
return self._name
|
|
|
|
def get_total_memory(self) -> int:
|
|
return self._total_memory
|
|
|
|
def __eq__(self, other):
|
|
if self is not None and other is not None and isinstance(self, CuPyDeviceInfo) and isinstance(other, CuPyDeviceInfo):
|
|
return self._index == other._index
|
|
return False
|
|
|
|
def __hash__(self):
|
|
return self._index
|
|
|
|
def __str__(self):
|
|
if self.is_cpu():
|
|
return "CPU"
|
|
else:
|
|
return f"[{self._index}] {self._name} [{(self._total_memory / 1024**3) :.3}Gb]"
|
|
|
|
def __repr__(self):
|
|
return f'{self.__class__.__name__} object: ' + self.__str__()
|
|
|
|
|
|
_cupy_devices = None
|
|
def get_available_devices() -> List[CuPyDeviceInfo]:
|
|
"""
|
|
returns a list of available CuPyDeviceInfo
|
|
"""
|
|
global _cupy_devices
|
|
if _cupy_devices is None:
|
|
devices = []
|
|
|
|
for i in range (cp.cuda.runtime.getDeviceCount()):
|
|
device_props = cp.cuda.runtime.getDeviceProperties(i)
|
|
devices.append ( CuPyDeviceInfo(index=i, name=device_props['name'].decode('utf-8'), total_memory=device_props['totalGlobalMem']))
|
|
|
|
_cupy_devices = devices
|
|
return _cupy_devices
|
|
|