dnn.py 文件源码

def _dnn_lib():
    if _dnn_lib.handle is None:
        import ctypes.util

        lib_name = ctypes.util.find_library('cudnn')
        if lib_name is None and sys.platform == 'win32':
            # Update these names when new versions of cudnn are supported.
            for name in ['cudnn64_5.dll', 'cudnn64_4.dll']:
                lib_name = ctypes.util.find_library(name)
                if lib_name:
                    break
        if lib_name is None:
            raise RuntimeError('Could not find cudnn library (looked for v4 and v5[.1])')
        _dnn_lib.handle = ctypes.cdll.LoadLibrary(lib_name)
        cudnn = _dnn_lib.handle
        cudnn.cudnnCreate.argtypes = [ctypes.POINTER(ctypes.c_void_p)]
        cudnn.cudnnCreate.restype = ctypes.c_int
        cudnn.cudnnDestroy.argtypes = [ctypes.c_void_p]
        cudnn.cudnnDestroy.restype = ctypes.c_int
    return _dnn_lib.handle