spark-nlp 2.6.3rc1__py2.py3-none-any.whl → 6.2.1__py2.py3-none-any.whl

sparknlp/training/_tf_graph_builders/tf2contrib/rnn.py

@@ -0,0 +1,245 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""RNN helpers for TensorFlow models."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import rnn
+from tensorflow.python.ops import variable_scope as vs
+
+
+def stack_bidirectional_rnn(cells_fw,
+                            cells_bw,
+                            inputs,
+                            initial_states_fw=None,
+                            initial_states_bw=None,
+                            dtype=None,
+                            sequence_length=None,
+                            scope=None):
+    """Creates a bidirectional recurrent neural network.
+
+    Stacks several bidirectional rnn layers. The combined forward and backward
+    layer outputs are used as input of the next layer. tf.bidirectional_rnn
+    does not allow to share forward and backward information between layers.
+    The input_size of the first forward and backward cells must match.
+    The initial state for both directions is zero and no intermediate states
+    are returned.
+
+    As described in https://arxiv.org/abs/1303.5778
+
+    Args:
+      cells_fw: List of instances of RNNCell, one per layer,
+        to be used for forward direction.
+      cells_bw: List of instances of RNNCell, one per layer,
+        to be used for backward direction.
+      inputs: A length T list of inputs, each a tensor of shape
+        [batch_size, input_size], or a nested tuple of such elements.
+      initial_states_fw: (optional) A list of the initial states (one per layer)
+        for the forward RNN.
+        Each tensor must has an appropriate type and shape
+        `[batch_size, cell_fw.state_size]`.
+      initial_states_bw: (optional) Same as for `initial_states_fw`, but using
+        the corresponding properties of `cells_bw`.
+      dtype: (optional) The data type for the initial state.  Required if
+        either of the initial states are not provided.
+      sequence_length: (optional) An int32/int64 vector, size `[batch_size]`,
+        containing the actual lengths for each of the sequences.
+      scope: VariableScope for the created subgraph; defaults to None.
+
+    Returns:
+      A tuple (outputs, output_state_fw, output_state_bw) where:
+        outputs is a length `T` list of outputs (one for each input), which
+          are depth-concatenated forward and backward outputs.
+        output_states_fw is the final states, one tensor per layer,
+          of the forward rnn.
+        output_states_bw is the final states, one tensor per layer,
+          of the backward rnn.
+
+    Raises:
+      TypeError: If `cell_fw` or `cell_bw` is not an instance of `RNNCell`.
+      ValueError: If inputs is None, not a list or an empty list.
+    """
+    if not cells_fw:
+        raise ValueError("Must specify at least one fw cell for BidirectionalRNN.")
+    if not cells_bw:
+        raise ValueError("Must specify at least one bw cell for BidirectionalRNN.")
+    if not isinstance(cells_fw, list):
+        raise ValueError("cells_fw must be a list of RNNCells (one per layer).")
+    if not isinstance(cells_bw, list):
+        raise ValueError("cells_bw must be a list of RNNCells (one per layer).")
+    if len(cells_fw) != len(cells_bw):
+        raise ValueError("Forward and Backward cells must have the same depth.")
+    if (initial_states_fw is not None and
+            (not isinstance(initial_states_fw, list) or
+             len(initial_states_fw) != len(cells_fw))):
+        raise ValueError(
+            "initial_states_fw must be a list of state tensors (one per layer).")
+    if (initial_states_bw is not None and
+            (not isinstance(initial_states_bw, list) or
+             len(initial_states_bw) != len(cells_bw))):
+        raise ValueError(
+            "initial_states_bw must be a list of state tensors (one per layer).")
+    states_fw = []
+    states_bw = []
+    prev_layer = inputs
+
+    with vs.variable_scope(scope or "stack_bidirectional_rnn"):
+        for i, (cell_fw, cell_bw) in enumerate(zip(cells_fw, cells_bw)):
+            initial_state_fw = None
+            initial_state_bw = None
+            if initial_states_fw:
+                initial_state_fw = initial_states_fw[i]
+            if initial_states_bw:
+                initial_state_bw = initial_states_bw[i]
+
+            with vs.variable_scope("cell_%d" % i) as cell_scope:
+                prev_layer, state_fw, state_bw = rnn.static_bidirectional_rnn(
+                    cell_fw,
+                    cell_bw,
+                    prev_layer,
+                    initial_state_fw=initial_state_fw,
+                    initial_state_bw=initial_state_bw,
+                    sequence_length=sequence_length,
+                    dtype=dtype,
+                    scope=cell_scope)
+            states_fw.append(state_fw)
+            states_bw.append(state_bw)
+
+    return prev_layer, tuple(states_fw), tuple(states_bw)
+
+
+def stack_bidirectional_dynamic_rnn(cells_fw,
+                                    cells_bw,
+                                    inputs,
+                                    initial_states_fw=None,
+                                    initial_states_bw=None,
+                                    dtype=None,
+                                    sequence_length=None,
+                                    parallel_iterations=None,
+                                    time_major=False,
+                                    scope=None,
+                                    swap_memory=False):
+    """Creates a dynamic bidirectional recurrent neural network.
+
+    Stacks several bidirectional rnn layers. The combined forward and backward
+    layer outputs are used as input of the next layer. tf.bidirectional_rnn
+    does not allow to share forward and backward information between layers.
+    The input_size of the first forward and backward cells must match.
+    The initial state for both directions is zero and no intermediate states
+    are returned.
+
+    Args:
+      cells_fw: List of instances of RNNCell, one per layer,
+        to be used for forward direction.
+      cells_bw: List of instances of RNNCell, one per layer,
+        to be used for backward direction.
+      inputs: The RNN inputs. this must be a tensor of shape:
+        `[batch_size, max_time, ...]`, or a nested tuple of such elements.
+      initial_states_fw: (optional) A list of the initial states (one per layer)
+        for the forward RNN.
+        Each tensor must has an appropriate type and shape
+        `[batch_size, cell_fw.state_size]`.
+      initial_states_bw: (optional) Same as for `initial_states_fw`, but using
+        the corresponding properties of `cells_bw`.
+      dtype: (optional) The data type for the initial state.  Required if
+        either of the initial states are not provided.
+      sequence_length: (optional) An int32/int64 vector, size `[batch_size]`,
+        containing the actual lengths for each of the sequences.
+      parallel_iterations: (Default: 32).  The number of iterations to run in
+        parallel.  Those operations which do not have any temporal dependency
+        and can be run in parallel, will be.  This parameter trades off
+        time for space.  Values >> 1 use more memory but take less time,
+        while smaller values use less memory but computations take longer.
+      time_major: The shape format of the inputs and outputs Tensors. If true,
+        these Tensors must be shaped [max_time, batch_size, depth]. If false,
+        these Tensors must be shaped [batch_size, max_time, depth]. Using
+        time_major = True is a bit more efficient because it avoids transposes at
+        the beginning and end of the RNN calculation. However, most TensorFlow
+        data is batch-major, so by default this function accepts input and emits
+        output in batch-major form.
+      scope: VariableScope for the created subgraph; defaults to None.
+      swap_memory: Transparently swap the tensors produced in forward inference
+        but needed for back prop from GPU to CPU.  This allows training RNNs
+        which would typically not fit on a single GPU, with very minimal (or no)
+        performance penalty.
+
+    Returns:
+      A tuple (outputs, output_state_fw, output_state_bw) where:
+        outputs: Output `Tensor` shaped:
+          `[batch_size, max_time, layers_output]`. Where layers_output
+          are depth-concatenated forward and backward outputs.
+        output_states_fw is the final states, one tensor per layer,
+          of the forward rnn.
+        output_states_bw is the final states, one tensor per layer,
+          of the backward rnn.
+
+    Raises:
+      TypeError: If `cell_fw` or `cell_bw` is not an instance of `RNNCell`.
+      ValueError: If inputs is `None`.
+    """
+    if not cells_fw:
+        raise ValueError("Must specify at least one fw cell for BidirectionalRNN.")
+    if not cells_bw:
+        raise ValueError("Must specify at least one bw cell for BidirectionalRNN.")
+    if not isinstance(cells_fw, list):
+        raise ValueError("cells_fw must be a list of RNNCells (one per layer).")
+    if not isinstance(cells_bw, list):
+        raise ValueError("cells_bw must be a list of RNNCells (one per layer).")
+    if len(cells_fw) != len(cells_bw):
+        raise ValueError("Forward and Backward cells must have the same depth.")
+    if (initial_states_fw is not None and
+            (not isinstance(initial_states_fw, list) or
+             len(initial_states_fw) != len(cells_fw))):
+        raise ValueError(
+            "initial_states_fw must be a list of state tensors (one per layer).")
+    if (initial_states_bw is not None and
+            (not isinstance(initial_states_bw, list) or
+             len(initial_states_bw) != len(cells_bw))):
+        raise ValueError(
+            "initial_states_bw must be a list of state tensors (one per layer).")
+
+    states_fw = []
+    states_bw = []
+    prev_layer = inputs
+
+    with vs.variable_scope(scope or "stack_bidirectional_rnn"):
+        for i, (cell_fw, cell_bw) in enumerate(zip(cells_fw, cells_bw)):
+            initial_state_fw = None
+            initial_state_bw = None
+            if initial_states_fw:
+                initial_state_fw = initial_states_fw[i]
+            if initial_states_bw:
+                initial_state_bw = initial_states_bw[i]
+
+            with vs.variable_scope("cell_%d" % i):
+                outputs, (state_fw, state_bw) = rnn.bidirectional_dynamic_rnn(
+                    cell_fw,
+                    cell_bw,
+                    prev_layer,
+                    initial_state_fw=initial_state_fw,
+                    initial_state_bw=initial_state_bw,
+                    sequence_length=sequence_length,
+                    parallel_iterations=parallel_iterations,
+                    dtype=dtype,
+                    swap_memory=swap_memory,
+                    time_major=time_major)
+                # Concat the outputs to create the new input.
+                prev_layer = array_ops.concat(outputs, 2)
+            states_fw.append(state_fw)
+            states_bw.append(state_bw)
+
+    return prev_layer, tuple(states_fw), tuple(states_bw)