Address GPU message in root dir

This commit is contained in:
Philip Kahn 2018-05-04 15:07:45 -07:00 committed by GitHub
parent 75295718e1
commit 51fbf01277
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions

View File

@ -97,26 +97,26 @@ if os.name == "nt":
if int(os.environ['CUDA_VISIBLE_DEVICES']) < 0: if int(os.environ['CUDA_VISIBLE_DEVICES']) < 0:
raise ValueError("ForceCPU") raise ValueError("ForceCPU")
try: try:
# Check a global
global DARKNET_FORCE_CPU global DARKNET_FORCE_CPU
if DARKNET_FORCE_CPU: if DARKNET_FORCE_CPU:
raise ValueError("ForceCPU") raise ValueError("ForceCPU")
except NameError: except NameError:
pass pass
#print(os.environ.keys()) # print(os.environ.keys())
#print("FORCE_CPU flag undefined, proceeding with GPU") # print("FORCE_CPU flag undefined, proceeding with GPU")
if not os.path.exists(winGPUdll): if not os.path.exists(winGPUdll):
raise ValueError("NoDLL") raise ValueError("NoDLL")
lib = CDLL(winGPUdll, RTLD_GLOBAL) lib = CDLL(winGPUdll, RTLD_GLOBAL)
except (KeyError, ValueError): except (KeyError, ValueError):
print("Notice: GPU-free mode")
hasGPU = False hasGPU = False
if os.path.exists(winNoGPUdll): if os.path.exists(winNoGPUdll):
lib = CDLL(winNoGPUdll, RTLD_GLOBAL) lib = CDLL(winNoGPUdll, RTLD_GLOBAL)
print("Notice: CPU-only mode")
else: else:
# Try the other way, in case no_gpu was # Try the other way, in case no_gpu was
# compile but not renamed # compile but not renamed
lib = CDLL(winGPUdll, RTLD_GLOBAL) lib = CDLL(winGPUdll, RTLD_GLOBAL)
print("Environment variables indicated a CPU run, but we didn't find `"+winNoGPUdll+"`. Trying a GPU run anyway.")
else: else:
lib = CDLL("./libdarknet.so", RTLD_GLOBAL) lib = CDLL("./libdarknet.so", RTLD_GLOBAL)
lib.network_width.argtypes = [c_void_p] lib.network_width.argtypes = [c_void_p]