EVOLUTION-MANAGER

Edit File: rnn_cell.py

# Lint as: python2, python3 # Copyright 2018 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. # ============================================================================== """TfLite BasicRnnCell wrapper. TODO(renjieliu): Find a better home for this one. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import itertools from tensorflow.lite.python.op_hint import OpHint from tensorflow.python.layers import base as base_layer from tensorflow.python.ops import array_ops from tensorflow.python.ops import clip_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops from tensorflow.python.ops import partitioned_variables from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util.tf_export import tf_export @tf_export(v1=["lite.experimental.nn.TfLiteRNNCell"]) class TfLiteRNNCell(rnn_cell_impl.LayerRNNCell): """The most basic RNN cell. This is used only for TfLite, it provides hints and it also makes the variables in the desired for the tflite ops. """ def __init__(self, num_units, activation=None, reuse=None, name=None, dtype=None, **kwargs): """Initializes the parameters for an RNN cell. Args: num_units: int, The number of units in the RNN cell. activation: Nonlinearity to use. Default: `tanh`. It could also be string that is within Keras activation function names. reuse: (optional) Python boolean describing whether to reuse variables in an existing scope. Raises an error if not `True` and the existing scope already has the given variables. name: String, the name of the layer. Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases. dtype: Default dtype of the layer (default of `None` means use the type of the first input). Required when `build` is called before `call`. **kwargs: Dict, keyword named properties for common layer attributes, like `trainable` etc when constructing the cell from configs of get_config(). Raises: ValueError: If the existing scope already has the given variables. """ super(TfLiteRNNCell, self).__init__( _reuse=reuse, name=name, dtype=dtype, **kwargs) # Inputs must be Rank-2. self.input_spec = base_layer.InputSpec(ndim=2) self._tflite_wrapper = OpHint("UnidirectionalSequenceRnn") self._num_units = num_units if activation: if activation != "tanh": raise ValueError("activation other than tanh is not supported") self._activation = math_ops.tanh else: self._activation = math_ops.tanh @property def state_size(self): return self._num_units @property def output_size(self): return self._num_units def build(self, inputs_shape): """Builds the RNN cell. Args: inputs_shape: Rnn input tensor shape. Raises: ValueError: If last dimension of the input shape is not known. """ if inputs_shape[-1] is None: raise ValueError("Expected inputs.shape[-1] to be known, saw shape: %s" % (inputs_shape,)) input_depth = inputs_shape[-1] def add_variable_wrapped(name, shape, initializer, index): var = self.add_weight(name, shape=shape, initializer=initializer) return self._tflite_wrapper.add_input( var, name=name, index_override=index) self._input_weights = add_variable_wrapped( "input_weights", [self._num_units, input_depth], None, 1) self._recurrent_weights = add_variable_wrapped( "recurrent_weights", [self._num_units, self._num_units], None, 2) self._bias = add_variable_wrapped( "bias", shape=[self._num_units], initializer=init_ops.zeros_initializer(dtype=self.dtype), index=3) self.built = True def call(self, inputs, state): """Most basic RNN: output = new_state = act(W * input + U * state + B).""" inputs = self._tflite_wrapper.add_input( inputs, tag="input", name="input", aggregate="stack", index_override=0) state = self._tflite_wrapper.add_input( state, tag="hidden_state", name="hidden_state", aggregate="first", index_override=4) weights = array_ops.transpose( array_ops.concat([self._input_weights, self._recurrent_weights], 1)) gate_inputs = math_ops.matmul(array_ops.concat([inputs, state], 1), weights) gate_inputs = nn_ops.bias_add(gate_inputs, self._bias) output = self._activation(gate_inputs) output = self._tflite_wrapper.add_output( output, tag="output", name="output", index_override=1, aggregate="stack") return output, output def get_config(self): config = { "num_units": self._num_units, "activation": "tanh", "reuse": self._reuse, } base_config = super(TfLiteRNNCell, self).get_config() return dict( itertools.chain(list(base_config.items()), list(config.items()))) @tf_export(v1=["lite.experimental.nn.TFLiteLSTMCell"]) class TFLiteLSTMCell(rnn_cell_impl.LayerRNNCell): """Long short-term memory unit (LSTM) recurrent network cell. This is used only for TfLite, it provides hints and it also makes the variables in the desired for the tflite ops (transposed and separated). The default non-peephole implementation is based on: https://pdfs.semanticscholar.org/1154/0131eae85b2e11d53df7f1360eeb6476e7f4.pdf Felix Gers, Jurgen Schmidhuber, and Fred Cummins. "Learning to forget: Continual prediction with LSTM." IET, 850-855, 1999. The peephole implementation is based on: https://research.google.com/pubs/archive/43905.pdf Hasim Sak, Andrew Senior, and Francoise Beaufays. "Long short-term memory recurrent neural network architectures for large scale acoustic modeling." INTERSPEECH, 2014. The class uses optional peep-hole connections, optional cell clipping, and an optional projection layer. Note that this cell is not optimized for performance. Please use `tf.contrib.cudnn_rnn.CudnnLSTM` for better performance on GPU, or `tf.contrib.rnn.LSTMBlockCell` and `tf.contrib.rnn.LSTMBlockFusedCell` for better performance on CPU. """ def __init__(self, num_units, use_peepholes=False, cell_clip=None, initializer=None, num_proj=None, proj_clip=None, num_unit_shards=None, num_proj_shards=None, forget_bias=1.0, state_is_tuple=True, activation=None, reuse=None, name=None, dtype=None): """Initialize the parameters for an LSTM cell. Args: num_units: int, The number of units in the LSTM cell. use_peepholes: bool, set True to enable diagonal/peephole connections. cell_clip: (optional) A float value, if provided the cell state is clipped by this value prior to the cell output activation. initializer: (optional) The initializer to use for the weight and projection matrices. num_proj: (optional) int, The output dimensionality for the projection matrices. If None, no projection is performed. proj_clip: (optional) A float value. If `num_proj > 0` and `proj_clip` is provided, then the projected values are clipped elementwise to within `[-proj_clip, proj_clip]`. num_unit_shards: Deprecated, will be removed by Jan. 2017. Use a variable_scope partitioner instead. num_proj_shards: Deprecated, will be removed by Jan. 2017. Use a variable_scope partitioner instead. forget_bias: Biases of the forget gate are initialized by default to 1 in order to reduce the scale of forgetting at the beginning of the training. Must set it manually to `0.0` when restoring from CudnnLSTM trained checkpoints. state_is_tuple: If True, accepted and returned states are 2-tuples of the `c_state` and `m_state`. If False, they are concatenated along the column axis. This latter behavior will soon be deprecated. activation: Activation function of the inner states. Default: `tanh`. reuse: (optional) Python boolean describing whether to reuse variables in an existing scope. If not `True`, and the existing scope already has the given variables, an error is raised. name: String, the name of the layer. Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases. dtype: Default dtype of the layer (default of `None` means use the type of the first input). Required when `build` is called before `call`. When restoring from CudnnLSTM-trained checkpoints, use `CudnnCompatibleLSTMCell` instead. """ super(TFLiteLSTMCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) # TODO(raziel): decide if we want to just support tuples (yes please!). if not state_is_tuple: logging.warn( "%s: Using a concatenated state is slower and will soon be " "deprecated. Use state_is_tuple=True.", self) if num_unit_shards is not None or num_proj_shards is not None: logging.warn( "%s: The num_unit_shards and proj_unit_shards parameters are " "deprecated and will be removed in Jan 2017. " "Use a variable scope with a partitioner instead.", self) # Inputs must be 2-dimensional. # TODO(raziel): layers stuff -- chop if un-layerizing Op. self.input_spec = base_layer.InputSpec(ndim=2) self._tflite_wrapper = OpHint("UnidirectionalSequenceLstm") self._num_units = num_units self._use_peepholes = use_peepholes self._cell_clip = cell_clip self._initializer = initializer self._num_proj = num_proj self._proj_clip = proj_clip self._num_unit_shards = num_unit_shards self._num_proj_shards = num_proj_shards self._forget_bias = forget_bias self._state_is_tuple = state_is_tuple if activation: if activation != "tanh": raise ValueError("activation other than tanh is not supported") self._activation = math_ops.tanh else: self._activation = math_ops.tanh self._output_size = num_proj if num_proj else num_units self._state_size = ( rnn_cell_impl.LSTMStateTuple(num_units, self._output_size) if state_is_tuple else num_units + self._output_size) @property def state_size(self): return self._state_size @property def output_size(self): return self._output_size def build(self, inputs_shape): """Build TfLite LSTM cell graph. Args: inputs_shape: The inputs_shape must be known, and is [batch_size, input_size] shape. Raises: ValueError: if the inputs_shape is invalid. """ if len(inputs_shape) != 2: raise ValueError( "inputs_shape must be 2-dimensional, saw shape: %s" % inputs_shape) input_depth = ( inputs_shape[1] if isinstance(inputs_shape[1], int) else inputs_shape[1].value) if input_depth is None: raise ValueError("Invalid inputs_shape, saw shape: %s" % inputs_shape) maybe_partitioner = ( partitioned_variables.fixed_size_partitioner(self._num_unit_shards) if self._num_unit_shards is not None else None) input_weight_shape = [self._num_units, input_depth] cell_weight_shape = [self._num_units, self._output_size] bias_shape = [self._num_units] def add_variable_wrapped(name, shape, initializer, index, partitioner): var = self.add_weight( name, shape=shape, initializer=initializer, partitioner=partitioner) return self._tflite_wrapper.add_input( var, name=name, index_override=index) weight_initializer = self._initializer if self.dtype is None: bias_initializer = init_ops.zeros_initializer else: bias_initializer = init_ops.zeros_initializer(dtype=self.dtype) forget_bias_initializer = init_ops.constant_initializer(self._forget_bias) self.input_to_input_w = add_variable_wrapped( "input_to_input_w", input_weight_shape, weight_initializer, 1, maybe_partitioner) self.input_to_forget_w = add_variable_wrapped( "input_to_forget_w", input_weight_shape, weight_initializer, 2, maybe_partitioner) self.input_to_cell_w = add_variable_wrapped( "input_to_cell_w", input_weight_shape, weight_initializer, 3, maybe_partitioner) self.input_to_output_w = add_variable_wrapped( "input_to_output_w", input_weight_shape, weight_initializer, 4, maybe_partitioner) self.cell_to_input_w = add_variable_wrapped( "cell_to_input_w", cell_weight_shape, weight_initializer, 5, maybe_partitioner) self.cell_to_forget_w = add_variable_wrapped( "cell_to_forget_w", cell_weight_shape, weight_initializer, 6, maybe_partitioner) self.cell_to_cell_w = add_variable_wrapped( "cell_to_cell_w", cell_weight_shape, weight_initializer, 7, maybe_partitioner) self.cell_to_output_w = add_variable_wrapped( "cell_to_output_w", cell_weight_shape, weight_initializer, 8, maybe_partitioner) self.input_bias = add_variable_wrapped( "input_bias", bias_shape, bias_initializer, 12, maybe_partitioner) self.forget_bias = add_variable_wrapped("forget_bias", bias_shape, forget_bias_initializer, 13, maybe_partitioner) self.cell_bias = add_variable_wrapped( "cell_bias", bias_shape, bias_initializer, 14, maybe_partitioner) self.output_bias = add_variable_wrapped( "output_bias", bias_shape, bias_initializer, 15, maybe_partitioner) # index 9, 10, 11. # f stands for forget, i stands for input and o stands for output. if self._use_peepholes: self._w_f_diag = add_variable_wrapped("w_f_diag", [self._num_units], self._initializer, 10, maybe_partitioner) self._w_i_diag = add_variable_wrapped("w_i_diag", [self._num_units], self._initializer, 9, maybe_partitioner) self._w_o_diag = add_variable_wrapped("w_o_diag", [self._num_units], self._initializer, 11, maybe_partitioner) # index 16 for proj kernel. if self._num_proj is not None: maybe_proj_partitioner = ( partitioned_variables.fixed_size_partitioner(self._num_proj_shards) if self._num_proj_shards is not None else None) self._proj_kernel = add_variable_wrapped( "projection/kernel", [self._num_proj, self._num_units], self._initializer, 16, partitioner=maybe_proj_partitioner) self.built = True def call(self, inputs, state): """Run one step of LSTM. Args: inputs: input Tensor, 2D, `[batch, num_units]`. state: if `state_is_tuple` is False, this must be a state Tensor, `2-D, [batch, state_size]`. If `state_is_tuple` is True, this must be a tuple of state Tensors, both `2-D`, with column sizes `c_state` and `m_state`. Returns: A tuple containing: - A `2-D, [batch, output_dim]`, Tensor representing the output of the LSTM after reading `inputs` when previous state was `state`. Here output_dim is: num_proj if num_proj was set, num_units otherwise. - Tensor(s) representing the new state of LSTM after reading `inputs` when the previous state was `state`. Same type and shape(s) as `state`. Raises: ValueError: If input size cannot be inferred from inputs via static shape inference. """ inputs = self._tflite_wrapper.add_input( inputs, tag="input", name="input", aggregate="stack", index_override=0) # Make sure inputs and bias_initializer has the same type. assert inputs.dtype == self.input_to_input_w.dtype num_proj = self._num_units if self._num_proj is None else self._num_proj sigmoid = math_ops.sigmoid if self._state_is_tuple: (c_prev, m_prev) = state else: c_prev = array_ops.slice(state, [0, 0], [-1, self._num_units]) m_prev = array_ops.slice(state, [0, self._num_units], [-1, num_proj]) # Note: For TfLite, cell_state is at index 19 while activation state at # index 18. c_prev = self._tflite_wrapper.add_input( c_prev, tag="c_prev", name="c_prev", aggregate="first", index_override=19) m_prev = self._tflite_wrapper.add_input( m_prev, tag="m_prev", name="m_prev", aggregate="first", index_override=18) input_size = inputs.shape.with_rank(2).dims[1] if input_size.value is None: raise ValueError("Could not infer input size from inputs.shape[-1]") inputs_and_m_prev = array_ops.concat([inputs, m_prev], axis=1) # i stands for input gate. # f stands for forget gate activation. # o outputs. # j output of LSTM unit. # c is the final state. # m is the output. i = nn_ops.bias_add( math_ops.matmul( inputs_and_m_prev, array_ops.concat([self.input_to_input_w, self.cell_to_input_w], axis=1), transpose_b=True), self.input_bias) f = nn_ops.bias_add( math_ops.matmul( inputs_and_m_prev, array_ops.concat([self.input_to_forget_w, self.cell_to_forget_w], axis=1), transpose_b=True), self.forget_bias) o = nn_ops.bias_add( math_ops.matmul( inputs_and_m_prev, array_ops.concat([self.input_to_output_w, self.cell_to_output_w], axis=1), transpose_b=True), self.output_bias) j = nn_ops.bias_add( math_ops.matmul( inputs_and_m_prev, array_ops.concat([self.input_to_cell_w, self.cell_to_cell_w], axis=1), transpose_b=True), self.cell_bias) # Diagonal connections if self._use_peepholes: c = ( sigmoid(f + self._w_f_diag * c_prev) * c_prev + sigmoid(i + self._w_i_diag * c_prev) * self._activation(j)) else: c = (sigmoid(f) * c_prev + sigmoid(i) * self._activation(j)) if self._cell_clip is not None: # pylint: disable=invalid-unary-operand-type c = clip_ops.clip_by_value(c, -self._cell_clip, self._cell_clip) # pylint: enable=invalid-unary-operand-type if self._use_peepholes: m = sigmoid(o + self._w_o_diag * c) * self._activation(c) else: m = sigmoid(o) * self._activation(c) if self._num_proj is not None: transposed_proj_kernel = array_ops.transpose(self._proj_kernel) m = math_ops.matmul(m, transposed_proj_kernel) if self._proj_clip is not None: # pylint: disable=invalid-unary-operand-type m = clip_ops.clip_by_value(m, -self._proj_clip, self._proj_clip) # pylint: enable=invalid-unary-operand-type c = self._tflite_wrapper.add_output( c, tag="c", name="c", aggregate="last", index_override=1) m = self._tflite_wrapper.add_output( m, tag="m", name="m", index_override=2, aggregate="stack") new_state = ( rnn_cell_impl.LSTMStateTuple(c, m) if self._state_is_tuple else array_ops.concat([c, m], 1)) return m, new_state def get_config(self): config = { "num_units": self._num_units, "use_peepholes": self._use_peepholes, "cell_clip": self._cell_clip, "num_proj": self._num_proj, "proj_clip": self._proj_clip, "num_unit_shards": self._num_unit_shards, "num_proj_shards": self._num_proj_shards, "forget_bias": self._forget_bias, "state_is_tuple": self._state_is_tuple, "activation": "tanh", "reuse": self._reuse, } base_config = super(TFLiteLSTMCell, self).get_config() return dict(list(base_config.items()) + list(config.items()))