Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow_addons/register.py: 32%

1import glob 

2from pathlib import Path 

3 

4import tensorflow as tf 

5 

6from tensorflow_addons.utils.resource_loader import get_path_to_datafile 

7 

8 

9def register_all(keras_objects: bool = True, custom_kernels: bool = True) -> None: 

10 """Register TensorFlow Addons' objects in TensorFlow global dictionaries. 

11 

12 When loading a Keras model that has a TF Addons' function, it is needed 

13 for this function to be known by the Keras deserialization process. 

14 

15 There are two ways to do this, either do 

16 

17 ```python 

18 tf.keras.models.load_model( 

19 "my_model.tf", 

20 custom_objects={"LAMB": tfa.image.optimizer.LAMB} 

21 ) 

22 ``` 

23 

24 or you can do: 

25 ```python 

26 tfa.register_all() 

27 tf.keras.models.load_model("my_model.tf") 

28 ``` 

29 

30 If the model contains custom ops (compiled ops) of TensorFlow Addons, 

31 and the graph is loaded with `tf.saved_model.load`, then custom ops need 

32 to be registered before to avoid an error of the type: 

33 

34 ``` 

35 tensorflow.python.framework.errors_impl.NotFoundError: Op type not registered 

36 '...' in binary running on ... Make sure the Op and Kernel are 

37 registered in the binary running in this process. 

38 ``` 

39 

40 In this case, the only way to make sure that the ops are registered is to call 

41 this function: 

42 

43 ```python 

44 tfa.register_all() 

45 tf.saved_model.load("my_model.tf") 

46 ``` 

47 

48 Note that you can call this function multiple times in the same process, 

49 it only has an effect the first time. Afterward, it's just a no-op. 

50 

51 Args: 

52 keras_objects: boolean, `True` by default. If `True`, register all 

53 Keras objects 

54 with `tf.keras.utils.register_keras_serializable(package="Addons")` 

55 If set to False, doesn't register any Keras objects 

56 of Addons in TensorFlow. 

57 custom_kernels: boolean, `True` by default. If `True`, loads all 

58 custom kernels of TensorFlow Addons with 

59 `tf.load_op_library("path/to/so/file.so")`. Loading the SO files 

60 register them automatically. If `False` doesn't load and register 

61 the shared objects files. Not that it might be useful to turn it off 

62 if your installation of Addons doesn't work well with custom ops. 

63 Returns: 

64 None 

65 """ 

66 if keras_objects: 

67 register_keras_objects() 

68 if custom_kernels: 

69 register_custom_kernels() 

70 

71 

72def register_keras_objects() -> None: 

73 # TODO: once layer_test is replaced by a public API 

74 # and we can used unregistered objects with it 

75 # we can remove all decorators. 

76 # And register Keras objects here. 

77 pass 

78 

79 

80def register_custom_kernels() -> None: 

81 all_shared_objects = _get_all_shared_objects() 

82 if not all_shared_objects: 

83 raise FileNotFoundError( 

84 "No shared objects files were found in the custom ops " 

85 "directory in Tensorflow Addons, check your installation again, " 

86 "or, if you don't need custom ops, call `tfa.register_all(custom_kernels=False)`" 

87 " instead." 

88 ) 

89 try: 

90 for shared_object in all_shared_objects: 

91 tf.load_op_library(shared_object) 

92 except tf.errors.NotFoundError as e: 

93 raise RuntimeError( 

94 "One of the shared objects ({}) could not be loaded. This may be " 

95 "due to a number of reasons (incompatible TensorFlow version, buiding from " 

96 "source with different flags, broken install of TensorFlow Addons...). If you " 

97 "wanted to register the shared objects because you needed them when loading your " 

98 "model, you should fix your install of TensorFlow Addons. If you don't " 

99 "use custom ops in your model, you can skip registering custom ops with " 

100 "`tfa.register_all(custom_kernels=False)`".format(shared_object) 

101 ) from e 

102 

103 

104def _get_all_shared_objects(): 

105 custom_ops_dir = get_path_to_datafile("custom_ops", is_so=True) 

106 all_shared_objects = glob.glob(custom_ops_dir + "/**/*.so", recursive=True) 

107 all_shared_objects = [x for x in all_shared_objects if Path(x).is_file()] 

108 return all_shared_objects