github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-07 13:32:09 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

update devicelib/nnlib to detect compute capability

Browse source
This commit is contained in:
iperov 2019-01-08 10:48:06 +04:00
parent 29c2375f5a
commit e8620919a7
5 changed files with 70 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

30
nnlib/devicelib.py
View file

@ -7,7 +7,9 @@ class devicelib:
force_gpu_idxs = None
choose_worst_gpu = False
gpu_idxs = []
gpu_names = []
gpu_total_vram_gb = 0
gpu_compute_caps = []
allow_growth = True
use_fp16 = False
cpu_only = False
@ -47,12 +49,16 @@ class devicelib:
else:
self.gpu_idxs = [gpu_idx]
if len(self.gpu_idxs) == 0:
self.cpu_only = True
else:
self.cpu_only = False
self.cpu_only = (len(self.gpu_idxs) == 0)
if not self.cpu_only:
self.gpu_total_vram_gb = devicelib.getDeviceVRAMTotalGb ( self.gpu_idxs[0] )
self.gpu_names = []
self.gpu_compute_caps = []
for gpu_idx in self.gpu_idxs:
self.gpu_names += [devicelib.getDeviceName(gpu_idx)]
self.gpu_compute_caps += [ devicelib.getDeviceComputeCapability ( gpu_idx ) ]
@staticmethod
def hasNVML():
try:
@ -206,4 +212,16 @@ class devicelib:
nvmlShutdown()
except:
pass
return result
return result
@staticmethod
def getDeviceComputeCapability(idx):
result = 0
try:
nvmlInit()
if idx < nvmlDeviceGetCount():
result = nvmlDeviceGetCudaComputeCapability(nvmlDeviceGetHandleByIndex(idx))
nvmlShutdown()
except:
pass
return result[0] * 10 + result[1]