DeepSpeech/util/shared_lib.py
2017-03-28 13:37:38 -04:00

47 lines
1.5 KiB
Python

from __future__ import print_function
from __future__ import absolute_import
from util.gpu import get_available_gpus
from ctypes import cdll
from sys import platform as _platform
def get_cupti_libname():
if _platform == 'linux' or _platform == 'linux2':
return 'libcupti.so'
elif _platform == 'darwin':
return 'libcupti.dylib'
elif _platform == 'win32':
return 'libcupti.dll'
def check_cupti():
# We want to ensure that user has properly configured its libs.
# We do this because dso load of libcupti will happen after a lot
# of computation happened, so easy to miss and loose time.
libname = get_cupti_libname()
cupti = check_so(libname)
if cupti is None:
print("INFO: No %s because no GPU, go ahead." % libname)
elif cupti is True:
print("INFO: Found %s." % libname)
else:
print("WARNING: Running on GPU but no %s could be found ; will be unable to report GPU VRAM usage." % libname)
def check_so(soname):
"""
Verify that we do have the 'soname' lib present in the system, and that it
can be loaded.
"""
if len(get_available_gpus()) == 0:
return None
# Try to force load lib, this would fail if the lib is not there :)
try:
lib = cdll.LoadLibrary(soname)
print("INFO: Found so as", lib)
assert lib.__class__.__name__ == 'CDLL'
assert lib._name == soname
return True
except OSError as ex:
print("WARNING:", ex)
return False