tensorflow/compiler/tests/lstm.py

# Copyright 2017 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.
# ==============================================================================
"""A simple LSTM layer with benchmarks.

This sets up a simple LSTM (Long Short Term Memory) layer, unrolled to a fixed
length sequence.  The only deviation from standard LSTM cells is that
activations are clipped, inspired by the GNMT machine translation model.
The GNMT paper has more details: https://arxiv.org/abs/1609.08144
"""

from six.moves import range

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variable_v1


def Clip(x):
  """Clips x to the range [-1., 1.]."""
  return math_ops.maximum(math_ops.minimum(x, 1.), -1.)


def LSTMCellWeightsShape(num_inputs, num_nodes):
  """Returns the shape of the weights for a single LSTM cell."""
  # Dimension 0 accounts for combining x with the previous m state.
  # Dimension 1 accounts for the in value and the (in, forget, out) gates.
  return [num_inputs + num_nodes, 4 * num_nodes]


def LSTMCell(weights, m_prev, c_prev, x, pad):
  """Unrolls a single LSTM cell with clipped activations forward by one step.

  Args:
    weights: Weight matrix with shape LSTMCellWeightsShape.
    m_prev: Previous m states with shape [batch_size, num_nodes].
    c_prev: Previous c states with shape [batch_size, num_nodes].
    x: Input with shape [batch_size, num_inputs].
    pad: Padding with shape [batch_size, 1].  Each padding value is either
        0 or 1, where 1 indicates padding; i.e. the input is shorter than the
        sequence length, and the (m, c) states should simply be passed through
        from the previous states.
  Returns:
    The next (m, c) states, each with shape [batch_size, num_nodes].
  """
  # Apply weights to the input and previous hidden state.
  # The matmul here is the "big" operation.
  xm = array_ops.concat([x, m_prev], 1)
  xmw = math_ops.matmul(xm, weights)

  # Element-wise ops for the standard LSTM cell, with clipped activations.
  # XLA can fuse these operations into a single loop.
  in_value, in_gate, forget_gate, out_gate = array_ops.split(
      value=xmw, num_or_size_splits=4, axis=1)
  in_value = math_ops.tanh(in_value)
  in_gate = math_ops.sigmoid(in_gate)
  forget_gate = math_ops.sigmoid(forget_gate)
  out_gate = math_ops.sigmoid(out_gate)
  c_next = Clip(Clip(forget_gate * c_prev) + Clip(in_gate * in_value))
  m_next = Clip(out_gate * c_next)

  # Account for padding.
  c_next = c_prev * pad + c_next * (1.0 - pad)
  m_next = m_prev * pad + m_next * (1.0 - pad)

  return m_next, c_next


def LSTMLayer(cell_name, weights, m, c, x_seq, pad_seq):
  """Unrolls a layer of LSTM cells forward by the sequence length.

  The sequence length is determined by the length of x_seq and pad_seq, which
  must be the same.

  Args:
    cell_name: Base name of each cell.
    weights: Weight matrix with shape LSTMCellWeightsShape.
    m: Initial m states with shape [batch_size, num_nodes].
    c: Initial c states with shape [batch_size, num_nodes].
    x_seq: List of inputs, each with shape [batch_size, num_inputs].
        The length of the list is the sequence length.
    pad_seq: List of paddings, each with shape [batch_size, 1].
        The length of the list is the sequence length.
        Each padding value is either 0 or 1, where 1 indicates padding;
        i.e. the input is shorter than the sequence length.
  Returns:
    List of per-sequence-step outputs, each with shape [batch_size, num_nodes].
  Raises:
    ValueError: If len(x_seq) != len(pad_seq).
  """
  if len(x_seq) != len(pad_seq):
    raise ValueError('length of x_seq(%d) != pad_seq(%d)' %
                     (len(x_seq), len(pad_seq)))
  out_seq = []
  for seq in range(len(x_seq)):
    with ops.name_scope('%s_%d' % (cell_name, seq)):
      m, c = LSTMCell(weights, m, c, x_seq[seq], pad_seq[seq])
      out_seq.append(array_ops.identity(m, name='out'))
  return out_seq


def RandomVar(shape, name=None):
  """Returns a variable of the given shape initialized to random values."""
  return variable_v1.VariableV1(
      random_ops.random_uniform(shape), dtype=dtypes.float32, name=name)


def RandomInputs(batch_size, seq_length, num_inputs):
  """Returns randomly initialized (x_seq, pad_seq) sequences."""
  x_seq = []
  pad_seq = []
  with ops.name_scope('inputs'):
    for seq in range(seq_length):
      x_seq.append(RandomVar([batch_size, num_inputs], name='x_seq_%d' % seq))
      # Real padding values are always a sequence of 0 followed by a
      # sequence of 1, but random values are fine for benchmarking.
      pad_seq.append(RandomVar([batch_size, 1], name='pad_seq_%d' % seq))
  return x_seq, pad_seq


def BuildLSTMLayer(batch_size, seq_length, num_inputs, num_nodes):
  """Builds a single LSTM layer with random weights and inputs.

  Args:
    batch_size: Inputs are fed in batches of this size.
    seq_length: The sequence length to unroll the LSTM layer.
    num_inputs: Dimension of inputs that are fed into each LSTM cell.
    num_nodes: The number of nodes in each LSTM cell.

  Returns:
    (out_seq, weights) pair.  The out_seq is a list of per-sequence-step
    outputs, each with shape [batch_size, num_nodes].  The weights are a list of
    weight variables that may be trained.
  """
  weights = RandomVar(
      LSTMCellWeightsShape(num_inputs, num_nodes), name='weights')
  m = array_ops.zeros([batch_size, num_nodes], name='init_m')
  c = array_ops.zeros([batch_size, num_nodes], name='init_c')
  x_seq, pad_seq = RandomInputs(batch_size, seq_length, num_inputs)

  out_seq = LSTMLayer('lstm', weights, m, c, x_seq, pad_seq)
  return out_seq, [weights]