Skip to content

Commit 9e13d30

Browse files
gabrieldemarmiesseseanpmorgan
gabrieldemarmiesse
authored and
seanpmorgan
committed
Lazy loading of shared libraries. (#869)
* Added lazy loading
1 parent 0a23fd2 commit 9e13d30

18 files changed

+77
-81
lines changed

‎tensorflow_addons/activations/gelu.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@
1818
from __future__ import print_function
1919

2020
import tensorflow as tf
21-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
21+
from tensorflow_addons.utils.resource_loader import LazySO
2222

23-
_activation_ops_so = tf.load_op_library(
24-
get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
23+
_activation_so = LazySO("custom_ops/activations/_activation_ops.so")
2524

2625

2726
@tf.keras.utils.register_keras_serializable(package='Addons')
@@ -44,10 +43,10 @@ def gelu(x, approximate=True):
4443
A `Tensor`. Has the same type as `x`.
4544
"""
4645
x = tf.convert_to_tensor(x)
47-
return _activation_ops_so.addons_gelu(x, approximate)
46+
return _activation_so.ops.addons_gelu(x, approximate)
4847

4948

5049
@tf.RegisterGradient("Addons>Gelu")
5150
def _gelu_grad(op, grad):
52-
return _activation_ops_so.addons_gelu_grad(grad, op.inputs[0],
51+
return _activation_so.ops.addons_gelu_grad(grad, op.inputs[0],
5352
op.get_attr("approximate"))

‎tensorflow_addons/activations/hardshrink.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@
1818
from __future__ import print_function
1919

2020
import tensorflow as tf
21-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
21+
from tensorflow_addons.utils.resource_loader import LazySO
2222

23-
_activation_ops_so = tf.load_op_library(
24-
get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
23+
_activation_so = LazySO("custom_ops/activations/_activation_ops.so")
2524

2625

2726
@tf.keras.utils.register_keras_serializable(package='Addons')
@@ -40,11 +39,11 @@ def hardshrink(x, lower=-0.5, upper=0.5):
4039
A `Tensor`. Has the same type as `x`.
4140
"""
4241
x = tf.convert_to_tensor(x)
43-
return _activation_ops_so.addons_hardshrink(x, lower, upper)
42+
return _activation_so.ops.addons_hardshrink(x, lower, upper)
4443

4544

4645
@tf.RegisterGradient("Addons>Hardshrink")
4746
def _hardshrink_grad(op, grad):
48-
return _activation_ops_so.addons_hardshrink_grad(grad, op.inputs[0],
47+
return _activation_so.ops.addons_hardshrink_grad(grad, op.inputs[0],
4948
op.get_attr("lower"),
5049
op.get_attr("upper"))

‎tensorflow_addons/activations/lisht.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@
1818
from __future__ import print_function
1919

2020
import tensorflow as tf
21-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
21+
from tensorflow_addons.utils.resource_loader import LazySO
2222

23-
_activation_ops_so = tf.load_op_library(
24-
get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
23+
_activation_so = LazySO("custom_ops/activations/_activation_ops.so")
2524

2625

2726
@tf.keras.utils.register_keras_serializable(package='Addons')
@@ -39,9 +38,9 @@ def lisht(x):
3938
A `Tensor`. Has the same type as `x`.
4039
"""
4140
x = tf.convert_to_tensor(x)
42-
return _activation_ops_so.addons_lisht(x)
41+
return _activation_so.ops.addons_lisht(x)
4342

4443

4544
@tf.RegisterGradient("Addons>Lisht")
4645
def _lisht_grad(op, grad):
47-
return _activation_ops_so.addons_lisht_grad(grad, op.inputs[0])
46+
return _activation_so.ops.addons_lisht_grad(grad, op.inputs[0])

‎tensorflow_addons/activations/mish.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@
1818
from __future__ import print_function
1919

2020
import tensorflow as tf
21-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
21+
from tensorflow_addons.utils.resource_loader import LazySO
2222

23-
_activation_ops_so = tf.load_op_library(
24-
get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
23+
_activation_so = LazySO("custom_ops/activations/_activation_ops.so")
2524

2625

2726
@tf.keras.utils.register_keras_serializable(package='Addons')
@@ -39,9 +38,9 @@ def mish(x):
3938
A `Tensor`. Has the same type as `x`.
4039
"""
4140
x = tf.convert_to_tensor(x)
42-
return _activation_ops_so.addons_mish(x)
41+
return _activation_so.ops.addons_mish(x)
4342

4443

4544
@tf.RegisterGradient("Addons>Mish")
4645
def _mish_grad(op, grad):
47-
return _activation_ops_so.addons_mish_grad(grad, op.inputs[0])
46+
return _activation_so.ops.addons_mish_grad(grad, op.inputs[0])

‎tensorflow_addons/activations/softshrink.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@
1818
from __future__ import print_function
1919

2020
import tensorflow as tf
21-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
21+
from tensorflow_addons.utils.resource_loader import LazySO
2222

23-
_activation_ops_so = tf.load_op_library(
24-
get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
23+
_activation_so = LazySO("custom_ops/activations/_activation_ops.so")
2524

2625

2726
@tf.keras.utils.register_keras_serializable(package='Addons')
@@ -40,11 +39,11 @@ def softshrink(x, lower=-0.5, upper=0.5):
4039
A `Tensor`. Has the same type as `x`.
4140
"""
4241
x = tf.convert_to_tensor(x)
43-
return _activation_ops_so.addons_softshrink(x, lower, upper)
42+
return _activation_so.ops.addons_softshrink(x, lower, upper)
4443

4544

4645
@tf.RegisterGradient("Addons>Softshrink")
4746
def _softshrink_grad(op, grad):
48-
return _activation_ops_so.addons_softshrink_grad(grad, op.inputs[0],
47+
return _activation_so.ops.addons_softshrink_grad(grad, op.inputs[0],
4948
op.get_attr("lower"),
5049
op.get_attr("upper"))

‎tensorflow_addons/activations/tanhshrink.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@
1818
from __future__ import print_function
1919

2020
import tensorflow as tf
21-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
21+
from tensorflow_addons.utils.resource_loader import LazySO
2222

23-
_activation_ops_so = tf.load_op_library(
24-
get_path_to_datafile("custom_ops/activations/_activation_ops.so"))
23+
_activation_so = LazySO("custom_ops/activations/_activation_ops.so")
2524

2625

2726
@tf.keras.utils.register_keras_serializable(package='Addons')
@@ -35,9 +34,9 @@ def tanhshrink(x):
3534
A `Tensor`. Has the same type as `features`.
3635
"""
3736
x = tf.convert_to_tensor(x)
38-
return _activation_ops_so.addons_tanhshrink(x)
37+
return _activation_so.ops.addons_tanhshrink(x)
3938

4039

4140
@tf.RegisterGradient("Addons>Tanhshrink")
4241
def _tanhshrink_grad(op, grad):
43-
return _activation_ops_so.addons_tanhshrink_grad(grad, op.inputs[0])
42+
return _activation_so.ops.addons_tanhshrink_grad(grad, op.inputs[0])

‎tensorflow_addons/image/connected_components.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,9 @@
2020

2121
import tensorflow as tf
2222

23-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
23+
from tensorflow_addons.utils.resource_loader import LazySO
2424

25-
_image_ops_so = tf.load_op_library(
26-
get_path_to_datafile("custom_ops/image/_image_ops.so"))
25+
_image_so = LazySO("custom_ops/image/_image_ops.so")
2726

2827

2928
@tf.function
@@ -62,7 +61,7 @@ def connected_components(images, name=None):
6261
raise TypeError(
6362
"images should have rank 2 (HW) or 3 (NHW). Static shape is %s"
6463
% image_or_images.get_shape())
65-
components = _image_ops_so.addons_image_connected_components(images)
64+
components = _image_so.ops.addons_image_connected_components(images)
6665

6766
# TODO(ringwalt): Component id renaming should be done in the op,
6867
# to avoid constructing multiple additional large tensors.

‎tensorflow_addons/image/distance_transform.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,9 @@
1919

2020
import tensorflow as tf
2121
from tensorflow_addons.image import utils as img_utils
22-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
22+
from tensorflow_addons.utils.resource_loader import LazySO
2323

24-
_image_ops_so = tf.load_op_library(
25-
get_path_to_datafile("custom_ops/image/_image_ops.so"))
24+
_image_so = LazySO("custom_ops/image/_image_ops.so")
2625

2726
tf.no_gradient("Addons>EuclideanDistanceTransform")
2827

@@ -64,6 +63,6 @@ def euclidean_dist_transform(images, dtype=tf.float32, name=None):
6463
raise TypeError("`dtype` must be float16, float32 or float64")
6564

6665
images = tf.cast(images, dtype)
67-
output = _image_ops_so.addons_euclidean_distance_transform(images)
66+
output = _image_so.ops.addons_euclidean_distance_transform(images)
6867

6968
return img_utils.from_4D_image(output, original_ndims)

‎tensorflow_addons/image/distort_image_ops.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@
1818
from __future__ import print_function
1919

2020
import tensorflow as tf
21-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
21+
from tensorflow_addons.utils.resource_loader import LazySO
2222

23-
_distort_image_ops = tf.load_op_library(
24-
get_path_to_datafile("custom_ops/image/_distort_image_ops.so"))
23+
_distort_image_so = LazySO("custom_ops/image/_distort_image_ops.so")
2524

2625

2726
# pylint: disable=invalid-name
@@ -141,7 +140,7 @@ def adjust_hsv_in_yiq(image,
141140
orig_dtype = image.dtype
142141
flt_image = tf.image.convert_image_dtype(image, tf.dtypes.float32)
143142

144-
rgb_altered = _distort_image_ops.addons_adjust_hsv_in_yiq(
143+
rgb_altered = _distort_image_so.ops.addons_adjust_hsv_in_yiq(
145144
flt_image, delta_hue, scale_saturation, scale_value)
146145

147146
return tf.image.convert_image_dtype(rgb_altered, orig_dtype)

‎tensorflow_addons/image/resampler_ops.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,9 @@
1919
from __future__ import print_function
2020

2121
import tensorflow as tf
22-
from tensorflow_addons.utils.resource_loader import get_path_to_datafile
22+
from tensorflow_addons.utils.resource_loader import LazySO
2323

24-
_resampler_ops = tf.load_op_library(
25-
get_path_to_datafile("custom_ops/image/_resampler_ops.so"))
24+
_resampler_so = LazySO("custom_ops/image/_resampler_ops.so")
2625

2726

2827
@tf.function
@@ -52,14 +51,15 @@ def resampler(data, warp, name=None):
5251
with tf.name_scope(name or "resampler"):
5352
data_tensor = tf.convert_to_tensor(data, name="data")
5453
warp_tensor = tf.convert_to_tensor(warp, name="warp")
55-
return _resampler_ops.addons_resampler(data_tensor, warp_tensor)
54+
return _resampler_so.ops.addons_resampler(data_tensor, warp_tensor)
5655

5756

5857
@tf.RegisterGradient("Addons>Resampler")
5958
def _resampler_grad(op, grad_output):
6059
data, warp = op.inputs
6160
grad_output_tensor = tf.convert_to_tensor(grad_output, name="grad_output")
62-
return _resampler_ops.addons_resampler_grad(data, warp, grad_output_tensor)
61+
return _resampler_so.ops.addons_resampler_grad(data, warp,
62+
grad_output_tensor)
6363

6464

6565
tf.no_gradient("Addons>ResamplerGrad")

0 commit comments

Comments
 (0)