提交 9483bc98 authored 作者: fo40225's avatar fo40225

change default dnn path on windows

上级 429a8fae
...@@ -323,23 +323,32 @@ def get_cuda_root(): ...@@ -323,23 +323,32 @@ def get_cuda_root():
return '' return ''
def default_dnn_path(suffix): def default_dnn_include_path():
def f(suffix=suffix): cuda_root = get_cuda_root()
cuda_root = get_cuda_root() if cuda_root == '':
if cuda_root == '': return ''
return '' return os.path.join(cuda_root, 'include')
return os.path.join(cuda_root, suffix)
return f
def default_dnn_library_path():
cuda_root = get_cuda_root()
if cuda_root == '':
return ''
if sys.platform == 'darwin':
return os.path.join(cuda_root, 'lib')
if sys.platform == 'win32':
return os.path.join(cuda_root, 'lib', 'x64')
return os.path.join(cuda_root, 'lib64')
AddConfigVar('dnn.include_path', AddConfigVar('dnn.include_path',
"Location of the cudnn header (defaults to the cuda root)", "Location of the cudnn header (defaults to the cuda root)",
StrParam(default_dnn_path('include')), StrParam(default_dnn_include_path()),
# Added elsewhere in the c key only when needed. # Added elsewhere in the c key only when needed.
in_c_key=False) in_c_key=False)
AddConfigVar('dnn.library_path', AddConfigVar('dnn.library_path',
"Location of the cudnn header (defaults to the cuda root)", "Location of the cudnn library (defaults to the cuda root)",
StrParam(default_dnn_path('lib' if sys.platform == 'darwin' else 'lib64')), StrParam(default_dnn_library_path()),
# Added elsewhere in the c key only when needed. # Added elsewhere in the c key only when needed.
in_c_key=False) in_c_key=False)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论