Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/keras/src/layers/rnn/gru_lstm_utils.py: 32%
84 statements
« prev ^ index » next coverage.py v7.4.0, created at 2024-01-03 07:57 +0000
1# Copyright 2019 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"""Utilities used by both the GRU and LSTM classes."""
18import uuid
20import tensorflow.compat.v2 as tf
22# isort: off
23from tensorflow.python.eager.context import get_device_name
25# The following string constants are used by Defun approach for unified backend
26# of LSTM and GRU.
27_FUNCTION_API_NAME_ATTRIBUTE = "api_implements"
28_FUNCTION_DEVICE_ATTRIBUTE = "api_preferred_device"
29CPU_DEVICE_NAME = "CPU"
30GPU_DEVICE_NAME = "GPU"
32# The following number constants are used to represent the runtime of the defun
33# backend function. Since the CPU/GPU implementation are mathematically same, we
34# need some signal for the function to indicate which function is executed. This
35# is for testing purpose to verify the correctness of swapping backend function.
36RUNTIME_UNKNOWN = 0
37RUNTIME_CPU = 1
38RUNTIME_GPU = 2
40CUDNN_AVAILABLE_MSG = "Layer %s will use cuDNN kernels when running on GPU."
41CUDNN_NOT_AVAILABLE_MSG = (
42 "Layer %s will not use cuDNN kernels since it "
43 "doesn't meet the criteria. It will "
44 "use a generic GPU kernel as fallback when running "
45 "on GPU."
46)
49def use_new_gru_lstm_impl():
50 return False
53# TODO(b/169707691): The wrapper can be removed if TFLite doesn't need to rely
54# on supportive attributes from LSTM/GRU.
55class DefunWrapper:
56 """A wrapper with no deep copy of the Defun in LSTM/GRU layer."""
58 def __init__(self, time_major, go_backwards, layer_name):
59 self.time_major = time_major
60 self.go_backwards = go_backwards
61 self.layer_name = layer_name
62 if self.layer_name not in ["lstm", "gru"]:
63 raise ValueError(
64 "Defun wrapper only applies to LSTM and GRU layer, "
65 "but given {}".format(self.layer_name)
66 )
67 # The first two attributes are added to support TFLite use case.
68 supportive_attributes = {
69 "time_major": self.time_major,
70 "go_backwards": self.go_backwards,
71 _FUNCTION_API_NAME_ATTRIBUTE: self.layer_name
72 + "_"
73 + str(uuid.uuid4()),
74 }
75 if self.layer_name == "lstm":
76 from keras.src.layers.rnn import (
77 lstm,
78 )
80 layer_func = lstm.lstm_with_backend_selection
81 else:
82 from keras.src.layers.rnn import (
83 gru,
84 )
86 layer_func = gru.gru_with_backend_selection
88 self.defun_layer = tf.function(
89 layer_func,
90 autograph=False,
91 experimental_attributes=supportive_attributes,
92 )
94 def __deepcopy__(self, memo):
95 new_wrapper = type(self)(
96 self.time_major, self.go_backwards, self.layer_name
97 )
98 memo[id(self)] = new_wrapper
99 return new_wrapper
102def canonical_to_params(weights, biases, shape, transpose_weights=False):
103 """Utility function convert variable to cuDNN compatible parameter.
105 Note that Keras weights for kernels are different from the cuDNN format.
106 Eg.:
108 ```
109 Keras cuDNN
110 [[0, 1, 2], <---> [[0, 2, 4],
111 [3, 4, 5]] [1, 3, 5]]
112 ```
114 If the input weights need to be in a unified format, then set
115 `transpose_weights=True` to convert the weights.
117 Args:
118 weights: list of weights for the individual kernels and recurrent kernels.
119 biases: list of biases for individual gate.
120 shape: the shape for the converted variables that will be feed to cuDNN.
121 transpose_weights: boolean, whether to transpose the weights.
123 Returns:
124 The converted weights that can be feed to cuDNN ops as param.
125 """
127 def convert(w):
128 return tf.transpose(w) if transpose_weights else w
130 weights = [tf.reshape(convert(x), shape) for x in weights]
131 biases = [tf.reshape(x, shape) for x in biases]
132 return tf.concat(weights + biases, axis=0)
135def is_sequence_right_padded(mask):
136 """Check the mask tensor and see if it right padded.
138 For cuDNN kernel, it uses the sequence length param to skip the tailing
139 timestep. If the data is left padded, or not a strict right padding (has
140 masked value in the middle of the sequence), then cuDNN kernel won't be work
141 properly in those cases.
143 Left padded data: [[False, False, True, True, True]].
144 Right padded data: [[True, True, True, False, False]].
145 Mixture of mask/unmasked data: [[True, False, True, False, False]].
147 Note that for the mixed data example above, the actually data RNN should see
148 are those 2 Trues (index 0 and 2), the index 1 False should be ignored and
149 not pollute the internal states.
151 Args:
152 mask: the Boolean tensor with shape [batch, timestep]
154 Returns:
155 boolean scalar tensor, whether the mask is strictly right padded.
156 """
157 max_seq_length = tf.shape(mask)[1]
158 count_of_true = tf.reduce_sum(tf.cast(mask, tf.int32), axis=1)
159 right_padded_mask = tf.sequence_mask(count_of_true, maxlen=max_seq_length)
160 return tf.reduce_all(tf.equal(mask, right_padded_mask))
163def has_fully_masked_sequence(mask):
164 # See https://github.com/tensorflow/tensorflow/issues/33148 for more
165 # details. Cudnn kernel will error out if the input sequence contains any
166 # fully masked data. We walk around this issue by rerouting the computation
167 # to standard kernel, until the issue on cudnn side has been fixed. For a
168 # fully masked sequence, it will contain all Falses. To make it easy to
169 # check, we inverse the boolean, check if any of the sequence has all True.
170 return tf.reduce_any(tf.reduce_all(tf.logical_not(mask), axis=1))
173def is_cudnn_supported_inputs(mask, time_major, sequence_lengths):
174 if tf.sysconfig.get_build_info()["is_rocm_build"]:
175 if (not time_major) and (sequence_lengths is not None):
176 return False
177 if mask is not None:
178 return tf.reduce_all(mask)
179 elif sequence_lengths is not None:
180 return tf.math.equal(
181 tf.reduce_min(sequence_lengths), tf.reduce_max(sequence_lengths)
182 )
183 else:
184 return True
185 if mask is None:
186 return True
187 if time_major:
188 mask = tf.transpose(mask)
190 return tf.logical_and(
191 is_sequence_right_padded(mask),
192 tf.logical_not(has_fully_masked_sequence(mask)),
193 )
196def calculate_sequence_by_mask(mask, time_major):
197 """Calculate the sequence length tensor (1-D) based on the masking tensor.
199 The masking tensor is a 2D boolean tensor with shape [batch, timestep]. For
200 any timestep that should be masked, the corresponding field will be False.
201 Consider the following example:
202 a = [[True, True, False, False],
203 [True, True, True, False]]
204 It is a (2, 4) tensor, and the corresponding sequence length result should
205 be 1D tensor with value [2, 3]. Note that the masking tensor must be right
206 padded that could be checked by, e.g., `is_sequence_right_padded()`.
208 Args:
209 mask: Boolean tensor with shape [batch, timestep] or [timestep, batch] if
210 time_major=True.
211 time_major: Boolean, which indicates whether the mask is time major or
212 batch major.
213 Returns:
214 sequence_length: 1D int32 tensor.
215 """
216 timestep_index = 0 if time_major else 1
217 return tf.reduce_sum(tf.cast(mask, tf.int32), axis=timestep_index)
220def generate_defun_backend(
221 unique_api_name, preferred_device, func, supportive_attributes
222):
223 function_attributes = {
224 _FUNCTION_API_NAME_ATTRIBUTE: unique_api_name,
225 _FUNCTION_DEVICE_ATTRIBUTE: preferred_device,
226 }
227 function_attributes.update(supportive_attributes)
228 return tf.function(
229 func, autograph=False, experimental_attributes=function_attributes
230 )
233def get_context_device_type():
234 """Parse the current context and return the device type, eg CPU/GPU."""
235 current_device = get_device_name()
236 if current_device is None:
237 return None
238 return tf.compat.v1.DeviceSpec.from_string(current_device).device_type
241def runtime(runtime_name):
242 with tf.device("/cpu:0"):
243 return tf.constant(runtime_name, dtype=tf.float32, name="runtime")
246def read_variable_value(v):
247 """Read the value of a variable if it is variable."""
248 if isinstance(v, tf.Variable):
249 return v.read_value()
250 return v
253def function_register(func, *args, **kwargs):
254 """Register a specialization of a `Function` into the graph.
256 This won't actually call the function with the inputs, and only put the
257 function definition into graph. Register function with different input param
258 will result into multiple version of functions registered in graph.
260 Args:
261 func: the `Function` instance that generated by a @defun
262 *args: input arguments for the Python function.
263 **kwargs: input keyword arguments for the Python function.
265 Returns:
266 a `ConcreteFunction` object specialized to inputs and execution context.
268 Raises:
269 ValueError: When the input function is not a defun wrapped python
270 function.
271 """
272 concrete_func = func.get_concrete_function(*args, **kwargs)
273 concrete_func.add_to_graph()
274 concrete_func.add_gradient_functions_to_graph()
275 return concrete_func