From 82fac052204570b91bbe0081ae17c65831f5a55a Mon Sep 17 00:00:00 2001 From: gabrieldemarmiesse Date: Sun, 12 Jan 2020 23:19:20 +0000 Subject: [PATCH 1/5] Added lazy loading --- tensorflow_addons/activations/softshrink.py | 13 ++++++------- tensorflow_addons/utils/resource_loader.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/tensorflow_addons/activations/softshrink.py b/tensorflow_addons/activations/softshrink.py index 5c285042de..06015f1bbf 100644 --- a/tensorflow_addons/activations/softshrink.py +++ b/tensorflow_addons/activations/softshrink.py @@ -18,10 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_activation_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/activations/_activation_ops.so")) +_activation_ops_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') @@ -40,11 +39,11 @@ def softshrink(x, lower=-0.5, upper=0.5): A `Tensor`. Has the same type as `x`. """ x = tf.convert_to_tensor(x) - return _activation_ops_so.addons_softshrink(x, lower, upper) + return _activation_ops_so.get().addons_softshrink(x, lower, upper) @tf.RegisterGradient("Addons>Softshrink") def _softshrink_grad(op, grad): - return _activation_ops_so.addons_softshrink_grad(grad, op.inputs[0], - op.get_attr("lower"), - op.get_attr("upper")) + return _activation_ops_so.get().addons_softshrink_grad(grad, op.inputs[0], + op.get_attr("lower"), + op.get_attr("upper")) diff --git a/tensorflow_addons/utils/resource_loader.py b/tensorflow_addons/utils/resource_loader.py index e518f1e2ae..dd8205d797 100644 --- a/tensorflow_addons/utils/resource_loader.py +++ b/tensorflow_addons/utils/resource_loader.py @@ -19,6 +19,8 @@ import os +import tensorflow as tf + def get_project_root(): """Returns project root folder.""" @@ -37,3 +39,14 @@ def get_path_to_datafile(path): """ root_dir = get_project_root() return os.path.join(root_dir, path.replace("/", os.sep)) + + +class LazyOpLoader: + def __init__(self, relative_path): + self.relative_path = relative_path + self.op = None + + def get(self): + if self.op is None: + self.op = tf.load_op_library(get_path_to_datafile(self.relative_path)) + return self.op From 1a9708614b2a7e03d0f495395614d90dd4e641e4 Mon Sep 17 00:00:00 2001 From: gabrieldemarmiesse Date: Mon, 13 Jan 2020 21:05:40 +0000 Subject: [PATCH 2/5] Used lazy loading everywhere possible. --- tensorflow_addons/activations/gelu.py | 9 ++++----- tensorflow_addons/activations/hardshrink.py | 9 ++++----- tensorflow_addons/activations/lisht.py | 9 ++++----- tensorflow_addons/activations/mish.py | 9 ++++----- tensorflow_addons/activations/softshrink.py | 10 +++++----- tensorflow_addons/activations/tanhshrink.py | 9 ++++----- tensorflow_addons/image/connected_components.py | 7 +++---- tensorflow_addons/image/distance_transform.py | 7 +++---- tensorflow_addons/image/distort_image_ops.py | 7 +++---- tensorflow_addons/image/resampler_ops.py | 10 +++++----- tensorflow_addons/image/transform_ops.py | 9 ++++----- tensorflow_addons/layers/optical_flow.py | 10 +++++----- tensorflow_addons/seq2seq/beam_search_decoder.py | 10 ++++------ .../seq2seq/beam_search_decoder_test.py | 8 +++----- tensorflow_addons/seq2seq/beam_search_ops_test.py | 14 ++++++-------- tensorflow_addons/text/parse_time_op.py | 9 ++++----- tensorflow_addons/text/skip_gram_ops.py | 7 +++---- tensorflow_addons/utils/resource_loader.py | 14 ++++++++------ 18 files changed, 76 insertions(+), 91 deletions(-) diff --git a/tensorflow_addons/activations/gelu.py b/tensorflow_addons/activations/gelu.py index 22c32127dc..37dd480899 100644 --- a/tensorflow_addons/activations/gelu.py +++ b/tensorflow_addons/activations/gelu.py @@ -18,10 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_activation_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/activations/_activation_ops.so")) +_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') @@ -44,10 +43,10 @@ def gelu(x, approximate=True): A `Tensor`. Has the same type as `x`. """ x = tf.convert_to_tensor(x) - return _activation_ops_so.addons_gelu(x, approximate) + return _activation_so.ops.addons_gelu(x, approximate) @tf.RegisterGradient("Addons>Gelu") def _gelu_grad(op, grad): - return _activation_ops_so.addons_gelu_grad(grad, op.inputs[0], + return _activation_so.ops.addons_gelu_grad(grad, op.inputs[0], op.get_attr("approximate")) diff --git a/tensorflow_addons/activations/hardshrink.py b/tensorflow_addons/activations/hardshrink.py index 3a7e36d9f9..9f488e79f4 100644 --- a/tensorflow_addons/activations/hardshrink.py +++ b/tensorflow_addons/activations/hardshrink.py @@ -18,10 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_activation_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/activations/_activation_ops.so")) +_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') @@ -40,11 +39,11 @@ def hardshrink(x, lower=-0.5, upper=0.5): A `Tensor`. Has the same type as `x`. """ x = tf.convert_to_tensor(x) - return _activation_ops_so.addons_hardshrink(x, lower, upper) + return _activation_so.ops.addons_hardshrink(x, lower, upper) @tf.RegisterGradient("Addons>Hardshrink") def _hardshrink_grad(op, grad): - return _activation_ops_so.addons_hardshrink_grad(grad, op.inputs[0], + return _activation_so.ops.addons_hardshrink_grad(grad, op.inputs[0], op.get_attr("lower"), op.get_attr("upper")) diff --git a/tensorflow_addons/activations/lisht.py b/tensorflow_addons/activations/lisht.py index 7f5473d0d9..8a04e0a19d 100644 --- a/tensorflow_addons/activations/lisht.py +++ b/tensorflow_addons/activations/lisht.py @@ -18,10 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_activation_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/activations/_activation_ops.so")) +_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') @@ -39,9 +38,9 @@ def lisht(x): A `Tensor`. Has the same type as `x`. """ x = tf.convert_to_tensor(x) - return _activation_ops_so.addons_lisht(x) + return _activation_so.ops.addons_lisht(x) @tf.RegisterGradient("Addons>Lisht") def _lisht_grad(op, grad): - return _activation_ops_so.addons_lisht_grad(grad, op.inputs[0]) + return _activation_so.ops.addons_lisht_grad(grad, op.inputs[0]) diff --git a/tensorflow_addons/activations/mish.py b/tensorflow_addons/activations/mish.py index b681be4bad..5f82f58e9b 100644 --- a/tensorflow_addons/activations/mish.py +++ b/tensorflow_addons/activations/mish.py @@ -18,10 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_activation_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/activations/_activation_ops.so")) +_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') @@ -39,9 +38,9 @@ def mish(x): A `Tensor`. Has the same type as `x`. """ x = tf.convert_to_tensor(x) - return _activation_ops_so.addons_mish(x) + return _activation_so.ops.addons_mish(x) @tf.RegisterGradient("Addons>Mish") def _mish_grad(op, grad): - return _activation_ops_so.addons_mish_grad(grad, op.inputs[0]) + return _activation_so.ops.addons_mish_grad(grad, op.inputs[0]) diff --git a/tensorflow_addons/activations/softshrink.py b/tensorflow_addons/activations/softshrink.py index 06015f1bbf..029a4db042 100644 --- a/tensorflow_addons/activations/softshrink.py +++ b/tensorflow_addons/activations/softshrink.py @@ -20,7 +20,7 @@ import tensorflow as tf from tensorflow_addons.utils.resource_loader import LazyOpLoader -_activation_ops_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") +_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') @@ -39,11 +39,11 @@ def softshrink(x, lower=-0.5, upper=0.5): A `Tensor`. Has the same type as `x`. """ x = tf.convert_to_tensor(x) - return _activation_ops_so.get().addons_softshrink(x, lower, upper) + return _activation_so.ops.addons_softshrink(x, lower, upper) @tf.RegisterGradient("Addons>Softshrink") def _softshrink_grad(op, grad): - return _activation_ops_so.get().addons_softshrink_grad(grad, op.inputs[0], - op.get_attr("lower"), - op.get_attr("upper")) + return _activation_so.ops.addons_softshrink_grad(grad, op.inputs[0], + op.get_attr("lower"), + op.get_attr("upper")) diff --git a/tensorflow_addons/activations/tanhshrink.py b/tensorflow_addons/activations/tanhshrink.py index e880262c0c..4548449604 100644 --- a/tensorflow_addons/activations/tanhshrink.py +++ b/tensorflow_addons/activations/tanhshrink.py @@ -18,10 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_activation_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/activations/_activation_ops.so")) +_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') @@ -35,9 +34,9 @@ def tanhshrink(x): A `Tensor`. Has the same type as `features`. """ x = tf.convert_to_tensor(x) - return _activation_ops_so.addons_tanhshrink(x) + return _activation_so.ops.addons_tanhshrink(x) @tf.RegisterGradient("Addons>Tanhshrink") def _tanhshrink_grad(op, grad): - return _activation_ops_so.addons_tanhshrink_grad(grad, op.inputs[0]) + return _activation_so.ops.addons_tanhshrink_grad(grad, op.inputs[0]) diff --git a/tensorflow_addons/image/connected_components.py b/tensorflow_addons/image/connected_components.py index 3b81985210..871b53d2ce 100644 --- a/tensorflow_addons/image/connected_components.py +++ b/tensorflow_addons/image/connected_components.py @@ -20,10 +20,9 @@ import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_image_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/image/_image_ops.so")) +_image_so = LazyOpLoader("custom_ops/image/_image_ops.so") @tf.function @@ -62,7 +61,7 @@ def connected_components(images, name=None): raise TypeError( "images should have rank 2 (HW) or 3 (NHW). Static shape is %s" % image_or_images.get_shape()) - components = _image_ops_so.addons_image_connected_components(images) + components = _image_so.ops.addons_image_connected_components(images) # TODO(ringwalt): Component id renaming should be done in the op, # to avoid constructing multiple additional large tensors. diff --git a/tensorflow_addons/image/distance_transform.py b/tensorflow_addons/image/distance_transform.py index 47e36c0318..8c63db3834 100644 --- a/tensorflow_addons/image/distance_transform.py +++ b/tensorflow_addons/image/distance_transform.py @@ -19,10 +19,9 @@ import tensorflow as tf from tensorflow_addons.image import utils as img_utils -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_image_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/image/_image_ops.so")) +_image_so = LazyOpLoader("custom_ops/image/_image_ops.so") tf.no_gradient("Addons>EuclideanDistanceTransform") @@ -64,6 +63,6 @@ def euclidean_dist_transform(images, dtype=tf.float32, name=None): raise TypeError("`dtype` must be float16, float32 or float64") images = tf.cast(images, dtype) - output = _image_ops_so.addons_euclidean_distance_transform(images) + output = _image_so.ops.addons_euclidean_distance_transform(images) return img_utils.from_4D_image(output, original_ndims) diff --git a/tensorflow_addons/image/distort_image_ops.py b/tensorflow_addons/image/distort_image_ops.py index 2a7f8dee1b..8325ca1dcd 100644 --- a/tensorflow_addons/image/distort_image_ops.py +++ b/tensorflow_addons/image/distort_image_ops.py @@ -18,10 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_distort_image_ops = tf.load_op_library( - get_path_to_datafile("custom_ops/image/_distort_image_ops.so")) +_distort_image_so = LazyOpLoader("custom_ops/image/_distort_image_ops.so") # pylint: disable=invalid-name @@ -141,7 +140,7 @@ def adjust_hsv_in_yiq(image, orig_dtype = image.dtype flt_image = tf.image.convert_image_dtype(image, tf.dtypes.float32) - rgb_altered = _distort_image_ops.addons_adjust_hsv_in_yiq( + rgb_altered = _distort_image_so.ops.addons_adjust_hsv_in_yiq( flt_image, delta_hue, scale_saturation, scale_value) return tf.image.convert_image_dtype(rgb_altered, orig_dtype) diff --git a/tensorflow_addons/image/resampler_ops.py b/tensorflow_addons/image/resampler_ops.py index 304bed5ff3..611ab71196 100644 --- a/tensorflow_addons/image/resampler_ops.py +++ b/tensorflow_addons/image/resampler_ops.py @@ -19,10 +19,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_resampler_ops = tf.load_op_library( - get_path_to_datafile("custom_ops/image/_resampler_ops.so")) +_resampler_so = LazyOpLoader("custom_ops/image/_resampler_ops.so") @tf.function @@ -52,14 +51,15 @@ def resampler(data, warp, name=None): with tf.name_scope(name or "resampler"): data_tensor = tf.convert_to_tensor(data, name="data") warp_tensor = tf.convert_to_tensor(warp, name="warp") - return _resampler_ops.addons_resampler(data_tensor, warp_tensor) + return _resampler_so.ops.addons_resampler(data_tensor, warp_tensor) @tf.RegisterGradient("Addons>Resampler") def _resampler_grad(op, grad_output): data, warp = op.inputs grad_output_tensor = tf.convert_to_tensor(grad_output, name="grad_output") - return _resampler_ops.addons_resampler_grad(data, warp, grad_output_tensor) + return _resampler_so.ops.addons_resampler_grad(data, warp, + grad_output_tensor) tf.no_gradient("Addons>ResamplerGrad") diff --git a/tensorflow_addons/image/transform_ops.py b/tensorflow_addons/image/transform_ops.py index 104856c893..08084a367e 100644 --- a/tensorflow_addons/image/transform_ops.py +++ b/tensorflow_addons/image/transform_ops.py @@ -19,10 +19,9 @@ import tensorflow as tf from tensorflow_addons.image import utils as img_utils -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_image_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/image/_image_ops.so")) +_image_so = LazyOpLoader("custom_ops/image/_image_ops.so") _IMAGE_DTYPES = set([ tf.dtypes.uint8, tf.dtypes.int32, tf.dtypes.int64, tf.dtypes.float16, @@ -98,7 +97,7 @@ def transform(images, "transforms should have rank 1 or 2, but got rank %d" % len( transforms.get_shape())) - output = _image_ops_so.addons_image_projective_transform_v2( + output = _image_so.ops.addons_image_projective_transform_v2( images, output_shape=output_shape, transforms=transforms, @@ -270,7 +269,7 @@ def _image_projective_transform_grad(op, grad): transforms = flat_transforms_to_matrices(transforms=transforms) inverse = tf.linalg.inv(transforms) transforms = matrices_to_flat_transforms(inverse) - output = _image_ops_so.addons_image_projective_transform_v2( + output = _image_so.ops.addons_image_projective_transform_v2( images=grad, transforms=transforms, output_shape=tf.shape(image_or_images)[1:3], diff --git a/tensorflow_addons/layers/optical_flow.py b/tensorflow_addons/layers/optical_flow.py index de831d2889..e59c297ccd 100644 --- a/tensorflow_addons/layers/optical_flow.py +++ b/tensorflow_addons/layers/optical_flow.py @@ -19,10 +19,10 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_correlation_cost_op_so = tf.load_op_library( - get_path_to_datafile("custom_ops/layers/_correlation_cost_ops.so")) +_correlation_cost_so = LazyOpLoader( + "custom_ops/layers/_correlation_cost_ops.so") def _correlation_cost(input_a, @@ -81,7 +81,7 @@ def _correlation_cost(input_a, """ with tf.name_scope(name or "correlation_cost"): - op_call = _correlation_cost_op_so.addons_correlation_cost + op_call = _correlation_cost_so.ops.addons_correlation_cost if data_format == "channels_last": op_data_format = "NHWC" @@ -120,7 +120,7 @@ def _correlation_cost_grad(op, grad_output): input_b = tf.convert_to_tensor(op.inputs[1], name="input_b") grad_output_tensor = tf.convert_to_tensor(grad_output, name="grad_output") - op_call = _correlation_cost_op_so.addons_correlation_cost_grad + op_call = _correlation_cost_so.ops.addons_correlation_cost_grad grads = op_call( input_a, input_b, diff --git a/tensorflow_addons/seq2seq/beam_search_decoder.py b/tensorflow_addons/seq2seq/beam_search_decoder.py index 85a34c5750..099ae59853 100644 --- a/tensorflow_addons/seq2seq/beam_search_decoder.py +++ b/tensorflow_addons/seq2seq/beam_search_decoder.py @@ -26,11 +26,9 @@ from tensorflow_addons.seq2seq import attention_wrapper from tensorflow_addons.seq2seq import decoder from tensorflow_addons.utils import keras_utils -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_beam_search_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/seq2seq/_beam_search_ops.so")) -gather_tree = _beam_search_ops_so.addons_gather_tree +_beam_search_so = LazyOpLoader("custom_ops/seq2seq/_beam_search_ops.so") class BeamSearchDecoderState( @@ -135,7 +133,7 @@ def gather_tree_from_array(t, parent_ids, sequence_length): max_sequence_lengths = tf.cast( tf.reduce_max(sequence_length, axis=1), tf.int32) - sorted_beam_ids = gather_tree( + sorted_beam_ids = _beam_search_so.ops.addons_gather_tree( step_ids=beam_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -342,7 +340,7 @@ def finalize(self, outputs, final_state, sequence_lengths): # Get max_sequence_length across all beams for each batch. max_sequence_lengths = tf.cast( tf.reduce_max(final_state.lengths, axis=1), tf.int32) - predicted_ids = gather_tree( + predicted_ids = _beam_search_so.ops.addons_gather_tree( outputs.predicted_ids, outputs.parent_ids, max_sequence_lengths=max_sequence_lengths, diff --git a/tensorflow_addons/seq2seq/beam_search_decoder_test.py b/tensorflow_addons/seq2seq/beam_search_decoder_test.py index 2730249fc7..fd27c9a47a 100644 --- a/tensorflow_addons/seq2seq/beam_search_decoder_test.py +++ b/tensorflow_addons/seq2seq/beam_search_decoder_test.py @@ -24,11 +24,9 @@ from tensorflow_addons.seq2seq import attention_wrapper from tensorflow_addons.seq2seq import beam_search_decoder from tensorflow_addons.utils import test_utils -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_beam_search_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/seq2seq/_beam_search_ops.so")) -gather_tree = _beam_search_ops_so.addons_gather_tree +_beam_search_so = LazyOpLoader("custom_ops/seq2seq/_beam_search_ops.so") class TestGatherTree(tf.test.TestCase): @@ -52,7 +50,7 @@ def test_gather_tree(self): [[2, 4, 4], [7, 6, 6], [8, 9, 10]]]).transpose([1, 0, 2]) - res = gather_tree( + res = _beam_search_so.ops.addons_gather_tree( predicted_ids, parent_ids, max_sequence_lengths=max_sequence_lengths, diff --git a/tensorflow_addons/seq2seq/beam_search_ops_test.py b/tensorflow_addons/seq2seq/beam_search_ops_test.py index dfc5d7a9dd..a64cdd4530 100644 --- a/tensorflow_addons/seq2seq/beam_search_ops_test.py +++ b/tensorflow_addons/seq2seq/beam_search_ops_test.py @@ -21,11 +21,9 @@ import numpy as np import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_beam_search_ops_so = tf.load_op_library( - get_path_to_datafile("custom_ops/seq2seq/_beam_search_ops.so")) -gather_tree = _beam_search_ops_so.addons_gather_tree +_beam_search_so = LazyOpLoader("custom_ops/seq2seq/_beam_search_ops.so") def _transpose_batch_time(x): @@ -43,7 +41,7 @@ def testGatherTreeOne(self): max_sequence_lengths = [3] expected_result = _transpose_batch_time([[[2, 2, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]) - beams = gather_tree( + beams = _beam_search_so.ops.addons_gather_tree( step_ids=step_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -63,7 +61,7 @@ def testBadParentValuesOnCPU(self): with tf.device("/cpu:0"): msg = r"parent id -1 at \(batch, time, beam\) == \(0, 0, 1\)" with self.assertRaisesOpError(msg): - beams = gather_tree( + beams = _beam_search_so.ops.addons_gather_tree( step_ids=step_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -86,7 +84,7 @@ def testBadParentValuesOnGPU(self): expected_result = _transpose_batch_time([[[2, -1, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]) with tf.device("/device:GPU:0"): - beams = gather_tree( + beams = _beam_search_so.ops.addons_gather_tree( step_ids=step_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -108,7 +106,7 @@ def testGatherTreeBatch(self): high=beam_width - 1, size=(max_time, batch_size, beam_width)) - beams = gather_tree( + beams = _beam_search_so.ops.addons_gather_tree( step_ids=step_ids.astype(np.int32), parent_ids=parent_ids.astype(np.int32), max_sequence_lengths=max_sequence_lengths, diff --git a/tensorflow_addons/text/parse_time_op.py b/tensorflow_addons/text/parse_time_op.py index 3369c83200..176d13ba00 100644 --- a/tensorflow_addons/text/parse_time_op.py +++ b/tensorflow_addons/text/parse_time_op.py @@ -19,10 +19,9 @@ import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_parse_time_op = tf.load_op_library( - get_path_to_datafile("custom_ops/text/_parse_time_op.so")) +_parse_time_so = LazyOpLoader("custom_ops/text/_parse_time_op.so") tf.no_gradient("Addons>ParseTime") @@ -82,5 +81,5 @@ def parse_time(time_string, time_format, output_unit): ValueError: If `output_unit` is not a valid value, if parsing `time_string` according to `time_format` failed. """ - return _parse_time_op.addons_parse_time(time_string, time_format, - output_unit) + return _parse_time_so.ops.addons_parse_time(time_string, time_format, + output_unit) diff --git a/tensorflow_addons/text/skip_gram_ops.py b/tensorflow_addons/text/skip_gram_ops.py index 412b9401ef..5cbe5c5d08 100644 --- a/tensorflow_addons/text/skip_gram_ops.py +++ b/tensorflow_addons/text/skip_gram_ops.py @@ -20,10 +20,9 @@ import csv import tensorflow as tf -from tensorflow_addons.utils.resource_loader import get_path_to_datafile +from tensorflow_addons.utils.resource_loader import LazyOpLoader -_skip_gram_ops = tf.load_op_library( - get_path_to_datafile("custom_ops/text/_skip_gram_ops.so")) +_skip_gram_so = LazyOpLoader("custom_ops/text/_skip_gram_ops.so") tf.no_gradient("Addons>SkipGramGenerateCandidates") @@ -170,7 +169,7 @@ def skip_gram_sample(input_tensor, seed=seed) seed1, seed2 = tf.compat.v1.get_seed(seed) - tokens, labels = _skip_gram_ops.addons_skip_gram_generate_candidates( + tokens, labels = _skip_gram_so.ops.addons_skip_gram_generate_candidates( input_tensor=input_tensor, min_skips=min_skips, max_skips=max_skips, diff --git a/tensorflow_addons/utils/resource_loader.py b/tensorflow_addons/utils/resource_loader.py index dd8205d797..ada09b9165 100644 --- a/tensorflow_addons/utils/resource_loader.py +++ b/tensorflow_addons/utils/resource_loader.py @@ -44,9 +44,11 @@ def get_path_to_datafile(path): class LazyOpLoader: def __init__(self, relative_path): self.relative_path = relative_path - self.op = None - - def get(self): - if self.op is None: - self.op = tf.load_op_library(get_path_to_datafile(self.relative_path)) - return self.op + self._ops = None + + @property + def ops(self): + if self._ops is None: + self._ops = tf.load_op_library( + get_path_to_datafile(self.relative_path)) + return self._ops From d19bb120afb922dc6b013928ed65c858f71270c8 Mon Sep 17 00:00:00 2001 From: gabrieldemarmiesse Date: Mon, 13 Jan 2020 21:08:04 +0000 Subject: [PATCH 3/5] Renamed class. --- tensorflow_addons/activations/gelu.py | 4 ++-- tensorflow_addons/activations/hardshrink.py | 4 ++-- tensorflow_addons/activations/lisht.py | 4 ++-- tensorflow_addons/activations/mish.py | 4 ++-- tensorflow_addons/activations/softshrink.py | 4 ++-- tensorflow_addons/activations/tanhshrink.py | 4 ++-- tensorflow_addons/image/connected_components.py | 4 ++-- tensorflow_addons/image/distance_transform.py | 4 ++-- tensorflow_addons/image/distort_image_ops.py | 4 ++-- tensorflow_addons/image/resampler_ops.py | 4 ++-- tensorflow_addons/image/transform_ops.py | 4 ++-- tensorflow_addons/layers/optical_flow.py | 4 ++-- tensorflow_addons/seq2seq/beam_search_decoder.py | 4 ++-- tensorflow_addons/seq2seq/beam_search_decoder_test.py | 4 ++-- tensorflow_addons/seq2seq/beam_search_ops_test.py | 4 ++-- tensorflow_addons/text/parse_time_op.py | 4 ++-- tensorflow_addons/text/skip_gram_ops.py | 4 ++-- tensorflow_addons/utils/resource_loader.py | 2 +- 18 files changed, 35 insertions(+), 35 deletions(-) diff --git a/tensorflow_addons/activations/gelu.py b/tensorflow_addons/activations/gelu.py index 37dd480899..f3f4d28a97 100644 --- a/tensorflow_addons/activations/gelu.py +++ b/tensorflow_addons/activations/gelu.py @@ -18,9 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") +_activation_so = LazySO("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') diff --git a/tensorflow_addons/activations/hardshrink.py b/tensorflow_addons/activations/hardshrink.py index 9f488e79f4..802391e725 100644 --- a/tensorflow_addons/activations/hardshrink.py +++ b/tensorflow_addons/activations/hardshrink.py @@ -18,9 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") +_activation_so = LazySO("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') diff --git a/tensorflow_addons/activations/lisht.py b/tensorflow_addons/activations/lisht.py index 8a04e0a19d..ab756a77d4 100644 --- a/tensorflow_addons/activations/lisht.py +++ b/tensorflow_addons/activations/lisht.py @@ -18,9 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") +_activation_so = LazySO("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') diff --git a/tensorflow_addons/activations/mish.py b/tensorflow_addons/activations/mish.py index 5f82f58e9b..4919598e5c 100644 --- a/tensorflow_addons/activations/mish.py +++ b/tensorflow_addons/activations/mish.py @@ -18,9 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") +_activation_so = LazySO("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') diff --git a/tensorflow_addons/activations/softshrink.py b/tensorflow_addons/activations/softshrink.py index 029a4db042..02af40931f 100644 --- a/tensorflow_addons/activations/softshrink.py +++ b/tensorflow_addons/activations/softshrink.py @@ -18,9 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") +_activation_so = LazySO("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') diff --git a/tensorflow_addons/activations/tanhshrink.py b/tensorflow_addons/activations/tanhshrink.py index 4548449604..4ec59b211f 100644 --- a/tensorflow_addons/activations/tanhshrink.py +++ b/tensorflow_addons/activations/tanhshrink.py @@ -18,9 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_activation_so = LazyOpLoader("custom_ops/activations/_activation_ops.so") +_activation_so = LazySO("custom_ops/activations/_activation_ops.so") @tf.keras.utils.register_keras_serializable(package='Addons') diff --git a/tensorflow_addons/image/connected_components.py b/tensorflow_addons/image/connected_components.py index 871b53d2ce..8f69d6e826 100644 --- a/tensorflow_addons/image/connected_components.py +++ b/tensorflow_addons/image/connected_components.py @@ -20,9 +20,9 @@ import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_image_so = LazyOpLoader("custom_ops/image/_image_ops.so") +_image_so = LazySO("custom_ops/image/_image_ops.so") @tf.function diff --git a/tensorflow_addons/image/distance_transform.py b/tensorflow_addons/image/distance_transform.py index 8c63db3834..f91642e627 100644 --- a/tensorflow_addons/image/distance_transform.py +++ b/tensorflow_addons/image/distance_transform.py @@ -19,9 +19,9 @@ import tensorflow as tf from tensorflow_addons.image import utils as img_utils -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_image_so = LazyOpLoader("custom_ops/image/_image_ops.so") +_image_so = LazySO("custom_ops/image/_image_ops.so") tf.no_gradient("Addons>EuclideanDistanceTransform") diff --git a/tensorflow_addons/image/distort_image_ops.py b/tensorflow_addons/image/distort_image_ops.py index 8325ca1dcd..7473f0d50c 100644 --- a/tensorflow_addons/image/distort_image_ops.py +++ b/tensorflow_addons/image/distort_image_ops.py @@ -18,9 +18,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_distort_image_so = LazyOpLoader("custom_ops/image/_distort_image_ops.so") +_distort_image_so = LazySO("custom_ops/image/_distort_image_ops.so") # pylint: disable=invalid-name diff --git a/tensorflow_addons/image/resampler_ops.py b/tensorflow_addons/image/resampler_ops.py index 611ab71196..c77aca7049 100644 --- a/tensorflow_addons/image/resampler_ops.py +++ b/tensorflow_addons/image/resampler_ops.py @@ -19,9 +19,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_resampler_so = LazyOpLoader("custom_ops/image/_resampler_ops.so") +_resampler_so = LazySO("custom_ops/image/_resampler_ops.so") @tf.function diff --git a/tensorflow_addons/image/transform_ops.py b/tensorflow_addons/image/transform_ops.py index 08084a367e..6dc60cc11a 100644 --- a/tensorflow_addons/image/transform_ops.py +++ b/tensorflow_addons/image/transform_ops.py @@ -19,9 +19,9 @@ import tensorflow as tf from tensorflow_addons.image import utils as img_utils -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_image_so = LazyOpLoader("custom_ops/image/_image_ops.so") +_image_so = LazySO("custom_ops/image/_image_ops.so") _IMAGE_DTYPES = set([ tf.dtypes.uint8, tf.dtypes.int32, tf.dtypes.int64, tf.dtypes.float16, diff --git a/tensorflow_addons/layers/optical_flow.py b/tensorflow_addons/layers/optical_flow.py index e59c297ccd..169278a831 100644 --- a/tensorflow_addons/layers/optical_flow.py +++ b/tensorflow_addons/layers/optical_flow.py @@ -19,9 +19,9 @@ from __future__ import print_function import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_correlation_cost_so = LazyOpLoader( +_correlation_cost_so = LazySO( "custom_ops/layers/_correlation_cost_ops.so") diff --git a/tensorflow_addons/seq2seq/beam_search_decoder.py b/tensorflow_addons/seq2seq/beam_search_decoder.py index 099ae59853..93b87d5683 100644 --- a/tensorflow_addons/seq2seq/beam_search_decoder.py +++ b/tensorflow_addons/seq2seq/beam_search_decoder.py @@ -26,9 +26,9 @@ from tensorflow_addons.seq2seq import attention_wrapper from tensorflow_addons.seq2seq import decoder from tensorflow_addons.utils import keras_utils -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_beam_search_so = LazyOpLoader("custom_ops/seq2seq/_beam_search_ops.so") +_beam_search_so = LazySO("custom_ops/seq2seq/_beam_search_ops.so") class BeamSearchDecoderState( diff --git a/tensorflow_addons/seq2seq/beam_search_decoder_test.py b/tensorflow_addons/seq2seq/beam_search_decoder_test.py index fd27c9a47a..3e37327783 100644 --- a/tensorflow_addons/seq2seq/beam_search_decoder_test.py +++ b/tensorflow_addons/seq2seq/beam_search_decoder_test.py @@ -24,9 +24,9 @@ from tensorflow_addons.seq2seq import attention_wrapper from tensorflow_addons.seq2seq import beam_search_decoder from tensorflow_addons.utils import test_utils -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_beam_search_so = LazyOpLoader("custom_ops/seq2seq/_beam_search_ops.so") +_beam_search_so = LazySO("custom_ops/seq2seq/_beam_search_ops.so") class TestGatherTree(tf.test.TestCase): diff --git a/tensorflow_addons/seq2seq/beam_search_ops_test.py b/tensorflow_addons/seq2seq/beam_search_ops_test.py index a64cdd4530..1453cb72fa 100644 --- a/tensorflow_addons/seq2seq/beam_search_ops_test.py +++ b/tensorflow_addons/seq2seq/beam_search_ops_test.py @@ -21,9 +21,9 @@ import numpy as np import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_beam_search_so = LazyOpLoader("custom_ops/seq2seq/_beam_search_ops.so") +_beam_search_so = LazySO("custom_ops/seq2seq/_beam_search_ops.so") def _transpose_batch_time(x): diff --git a/tensorflow_addons/text/parse_time_op.py b/tensorflow_addons/text/parse_time_op.py index 176d13ba00..a389cdcf76 100644 --- a/tensorflow_addons/text/parse_time_op.py +++ b/tensorflow_addons/text/parse_time_op.py @@ -19,9 +19,9 @@ import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_parse_time_so = LazyOpLoader("custom_ops/text/_parse_time_op.so") +_parse_time_so = LazySO("custom_ops/text/_parse_time_op.so") tf.no_gradient("Addons>ParseTime") diff --git a/tensorflow_addons/text/skip_gram_ops.py b/tensorflow_addons/text/skip_gram_ops.py index 5cbe5c5d08..899e445201 100644 --- a/tensorflow_addons/text/skip_gram_ops.py +++ b/tensorflow_addons/text/skip_gram_ops.py @@ -20,9 +20,9 @@ import csv import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazyOpLoader +from tensorflow_addons.utils.resource_loader import LazySO -_skip_gram_so = LazyOpLoader("custom_ops/text/_skip_gram_ops.so") +_skip_gram_so = LazySO("custom_ops/text/_skip_gram_ops.so") tf.no_gradient("Addons>SkipGramGenerateCandidates") diff --git a/tensorflow_addons/utils/resource_loader.py b/tensorflow_addons/utils/resource_loader.py index ada09b9165..3a08066206 100644 --- a/tensorflow_addons/utils/resource_loader.py +++ b/tensorflow_addons/utils/resource_loader.py @@ -41,7 +41,7 @@ def get_path_to_datafile(path): return os.path.join(root_dir, path.replace("/", os.sep)) -class LazyOpLoader: +class LazySO: def __init__(self, relative_path): self.relative_path = relative_path self._ops = None From 269041d9c1409c0a9344f038186ddd2c5f39b4f5 Mon Sep 17 00:00:00 2001 From: gabrieldemarmiesse Date: Mon, 13 Jan 2020 21:35:33 +0000 Subject: [PATCH 4/5] Now the import should work. --- tensorflow_addons/seq2seq/beam_search_decoder.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tensorflow_addons/seq2seq/beam_search_decoder.py b/tensorflow_addons/seq2seq/beam_search_decoder.py index 93b87d5683..d1cbbb00ac 100644 --- a/tensorflow_addons/seq2seq/beam_search_decoder.py +++ b/tensorflow_addons/seq2seq/beam_search_decoder.py @@ -31,6 +31,10 @@ _beam_search_so = LazySO("custom_ops/seq2seq/_beam_search_ops.so") +def gather_tree(*args, **kwargs): + return _beam_search_so.ops.addons_gather_tree(*args, **kwargs) + + class BeamSearchDecoderState( collections.namedtuple("BeamSearchDecoderState", ("cell_state", "log_probs", "finished", @@ -133,7 +137,7 @@ def gather_tree_from_array(t, parent_ids, sequence_length): max_sequence_lengths = tf.cast( tf.reduce_max(sequence_length, axis=1), tf.int32) - sorted_beam_ids = _beam_search_so.ops.addons_gather_tree( + sorted_beam_ids = gather_tree( step_ids=beam_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -340,7 +344,7 @@ def finalize(self, outputs, final_state, sequence_lengths): # Get max_sequence_length across all beams for each batch. max_sequence_lengths = tf.cast( tf.reduce_max(final_state.lengths, axis=1), tf.int32) - predicted_ids = _beam_search_so.ops.addons_gather_tree( + predicted_ids = gather_tree( outputs.predicted_ids, outputs.parent_ids, max_sequence_lengths=max_sequence_lengths, From e18db1a443e847412eb02af17ea43da532227c39 Mon Sep 17 00:00:00 2001 From: gabrieldemarmiesse Date: Mon, 13 Jan 2020 21:44:44 +0000 Subject: [PATCH 5/5] Imported function from file. --- .../seq2seq/beam_search_decoder_test.py | 7 ++----- tensorflow_addons/seq2seq/beam_search_ops_test.py | 12 +++++------- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/tensorflow_addons/seq2seq/beam_search_decoder_test.py b/tensorflow_addons/seq2seq/beam_search_decoder_test.py index 3e37327783..f12fb6c344 100644 --- a/tensorflow_addons/seq2seq/beam_search_decoder_test.py +++ b/tensorflow_addons/seq2seq/beam_search_decoder_test.py @@ -22,11 +22,8 @@ import tensorflow as tf from tensorflow_addons.seq2seq import attention_wrapper -from tensorflow_addons.seq2seq import beam_search_decoder +from tensorflow_addons.seq2seq import beam_search_decoder, gather_tree from tensorflow_addons.utils import test_utils -from tensorflow_addons.utils.resource_loader import LazySO - -_beam_search_so = LazySO("custom_ops/seq2seq/_beam_search_ops.so") class TestGatherTree(tf.test.TestCase): @@ -50,7 +47,7 @@ def test_gather_tree(self): [[2, 4, 4], [7, 6, 6], [8, 9, 10]]]).transpose([1, 0, 2]) - res = _beam_search_so.ops.addons_gather_tree( + res = gather_tree( predicted_ids, parent_ids, max_sequence_lengths=max_sequence_lengths, diff --git a/tensorflow_addons/seq2seq/beam_search_ops_test.py b/tensorflow_addons/seq2seq/beam_search_ops_test.py index 1453cb72fa..789ed72183 100644 --- a/tensorflow_addons/seq2seq/beam_search_ops_test.py +++ b/tensorflow_addons/seq2seq/beam_search_ops_test.py @@ -21,9 +21,7 @@ import numpy as np import tensorflow as tf -from tensorflow_addons.utils.resource_loader import LazySO - -_beam_search_so = LazySO("custom_ops/seq2seq/_beam_search_ops.so") +from tensorflow_addons.seq2seq import gather_tree def _transpose_batch_time(x): @@ -41,7 +39,7 @@ def testGatherTreeOne(self): max_sequence_lengths = [3] expected_result = _transpose_batch_time([[[2, 2, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]) - beams = _beam_search_so.ops.addons_gather_tree( + beams = gather_tree( step_ids=step_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -61,7 +59,7 @@ def testBadParentValuesOnCPU(self): with tf.device("/cpu:0"): msg = r"parent id -1 at \(batch, time, beam\) == \(0, 0, 1\)" with self.assertRaisesOpError(msg): - beams = _beam_search_so.ops.addons_gather_tree( + beams = gather_tree( step_ids=step_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -84,7 +82,7 @@ def testBadParentValuesOnGPU(self): expected_result = _transpose_batch_time([[[2, -1, 2], [6, 5, 6], [7, 8, 9], [10, 10, 10]]]) with tf.device("/device:GPU:0"): - beams = _beam_search_so.ops.addons_gather_tree( + beams = gather_tree( step_ids=step_ids, parent_ids=parent_ids, max_sequence_lengths=max_sequence_lengths, @@ -106,7 +104,7 @@ def testGatherTreeBatch(self): high=beam_width - 1, size=(max_time, batch_size, beam_width)) - beams = _beam_search_so.ops.addons_gather_tree( + beams = gather_tree( step_ids=step_ids.astype(np.int32), parent_ids=parent_ids.astype(np.int32), max_sequence_lengths=max_sequence_lengths,