github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-07 05:22:06 -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

32
nnlib/nnlib.py
View file

@ -9,7 +9,7 @@ from .devicelib import devicelib
class nnlib(object):
device = devicelib #forwards nnlib.devicelib to device in order to use nnlib as standalone lib
DeviceConfig = devicelib.Config
prefer_DeviceConfig = DeviceConfig() #default is one best GPU
active_DeviceConfig = DeviceConfig() #default is one best GPU
dlib = None
keras = None
@ -125,12 +125,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
def import_tf(device_config = None):
if nnlib.tf is not None:
return nnlib.code_import_tf
if device_config is None:
device_config = nnlib.prefer_DeviceConfig
else:
nnlib.prefer_DeviceConfig = device_config
if 'TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1':
suppressor = std_utils.suppress_stdout_stderr().__enter__()
else:
@ -144,6 +139,26 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
import tensorflow as tf
nnlib.tf = tf
if device_config is None:
device_config = nnlib.active_DeviceConfig
tf_ver = [int(x) for x in tf.VERSION.split('.')]
req_cap = 35
if tf_ver[0] > 1 or (tf_ver[0] == 1 and tf_ver[1] >= 11):
req_cap = 37
if not device_config.cpu_only and device_config.gpu_compute_caps[0] < req_cap:
if suppressor is not None:
suppressor.__exit__()
print ("%s does not meet minimum required compute capability: %d.%d. Falling back to CPU mode." % ( device_config.gpu_names[0], req_cap // 10, req_cap % 10 ) )
device_config = nnlib.DeviceConfig(cpu_only=True)
if suppressor is not None:
suppressor.__enter__()
nnlib.active_DeviceConfig = device_config
if device_config.cpu_only:
config = tf.ConfigProto( device_count = {'GPU': 0} )
else:
@ -160,6 +175,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
if suppressor is not None:
suppressor.__exit__()
nnlib.__initialize_tf_functions()
nnlib.code_import_tf = compile (nnlib.code_import_tf_string,'','exec')
@ -367,7 +383,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
return nnlib.code_import_keras
nnlib.import_tf(device_config)
device_config = nnlib.prefer_DeviceConfig
device_config = nnlib.active_DeviceConfig
if 'TF_SUPPRESS_STD' in os.environ.keys() and os.environ['TF_SUPPRESS_STD'] == '1':
suppressor = std_utils.suppress_stdout_stderr().__enter__()