1# Copyright 2015 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15 16"""Function for loading TensorFlow plugins.""" 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import errno 22import hashlib 23import imp 24import os 25import platform 26import sys 27import threading # pylint: disable=unused-import 28 29from tensorflow.core.framework import op_def_pb2 30from tensorflow.core.lib.core import error_codes_pb2 # pylint: disable=unused-import 31from tensorflow.python import pywrap_tensorflow as py_tf 32from tensorflow.python.lib.io import file_io 33from tensorflow.python.util import compat 34from tensorflow.python.util import deprecation 35from tensorflow.python.util.tf_export import tf_export 36 37 38@tf_export('load_op_library') 39def load_op_library(library_filename): 40 """Loads a TensorFlow plugin, containing custom ops and kernels. 41 42 Pass "library_filename" to a platform-specific mechanism for dynamically 43 loading a library. The rules for determining the exact location of the 44 library are platform-specific and are not documented here. When the 45 library is loaded, ops and kernels registered in the library via the 46 `REGISTER_*` macros are made available in the TensorFlow process. Note 47 that ops with the same name as an existing op are rejected and not 48 registered with the process. 49 50 Args: 51 library_filename: Path to the plugin. 52 Relative or absolute filesystem path to a dynamic library file. 53 54 Returns: 55 A python module containing the Python wrappers for Ops defined in 56 the plugin. 57 58 Raises: 59 RuntimeError: when unable to load the library or get the python wrappers. 60 """ 61 lib_handle = py_tf.TF_LoadLibrary(library_filename) 62 63 op_list_str = py_tf.TF_GetOpList(lib_handle) 64 op_list = op_def_pb2.OpList() 65 op_list.ParseFromString(compat.as_bytes(op_list_str)) 66 wrappers = py_tf.GetPythonWrappers(op_list_str) 67 68 # Delete the library handle to release any memory held in C 69 # that are no longer needed. 70 py_tf.TF_DeleteLibraryHandle(lib_handle) 71 72 # Get a unique name for the module. 73 module_name = hashlib.md5(wrappers).hexdigest() 74 if module_name in sys.modules: 75 return sys.modules[module_name] 76 module = imp.new_module(module_name) 77 # pylint: disable=exec-used 78 exec(wrappers, module.__dict__) 79 # Stash away the library handle for making calls into the dynamic library. 80 module.LIB_HANDLE = lib_handle 81 # OpDefs of the list of ops defined in the library. 82 module.OP_LIST = op_list 83 sys.modules[module_name] = module 84 return module 85 86 87@deprecation.deprecated(date=None, 88 instructions='Use `tf.load_library` instead.') 89@tf_export(v1=['load_file_system_library']) 90def load_file_system_library(library_filename): 91 """Loads a TensorFlow plugin, containing file system implementation. 92 93 Pass `library_filename` to a platform-specific mechanism for dynamically 94 loading a library. The rules for determining the exact location of the 95 library are platform-specific and are not documented here. 96 97 Args: 98 library_filename: Path to the plugin. 99 Relative or absolute filesystem path to a dynamic library file. 100 101 Returns: 102 None. 103 104 Raises: 105 RuntimeError: when unable to load the library. 106 """ 107 py_tf.TF_LoadLibrary(library_filename) 108 109 110def _is_shared_object(filename): 111 """Check the file to see if it is a shared object, only using extension.""" 112 if platform.system() == 'Linux': 113 if filename.endswith('.so'): 114 return True 115 else: 116 index = filename.rfind('.so.') 117 if index == -1: 118 return False 119 else: 120 # A shared object with the API version in filename 121 return filename[index + 4].isdecimal() 122 elif platform.system() == 'Darwin': 123 return filename.endswith('.dylib') 124 elif platform.system() == 'Windows': 125 return filename.endswith('.dll') 126 else: 127 return False 128 129 130@tf_export('load_library') 131def load_library(library_location): 132 """Loads a TensorFlow plugin. 133 134 "library_location" can be a path to a specific shared object, or a folder. 135 If it is a folder, all sahred objects that are named "libtfkernel*" will be 136 loaded. When the library is loaded, kernels registered in the library via the 137 `REGISTER_*` macros are made available in the TensorFlow process. 138 139 Args: 140 library_location: Path to the plugin or the folder of plugins. 141 Relative or absolute filesystem path to a dynamic library file or folder. 142 143 Returns: 144 None 145 146 Raises: 147 OSError: When the file to be loaded is not found. 148 RuntimeError: when unable to load the library. 149 """ 150 if file_io.file_exists(library_location): 151 if file_io.is_directory(library_location): 152 directory_contents = file_io.list_directory(library_location) 153 154 kernel_libraries = [ 155 os.path.join(library_location, f) for f in directory_contents 156 if _is_shared_object(f)] 157 else: 158 kernel_libraries = [library_location] 159 160 for lib in kernel_libraries: 161 py_tf.TF_LoadLibrary(lib) 162 163 else: 164 raise OSError( 165 errno.ENOENT, 166 'The file or folder to load kernel libraries from does not exist.', 167 library_location) 168 169