deeppavlov/core/layers/tf_layers.py
def cudnn_compatible_gru(units, n_hidden, n_layers=1, trainable_initial_states=False,
                         seq_lengths=None, input_initial_h=None, name='cudnn_gru', reuse=False):
    """ CuDNN Compatible GRU implementation.
        It should be used to load models saved with CudnnGRUCell to run on CPU.

        Args:
            units: tf.Tensor with dimensions [B x T x F], where
                B - batch size
                T - number of tokens
                F - features

            n_hidden: dimensionality of hidden state
            trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
            seq_lengths: tensor of sequence lengths with dimension [B]
            n_layers: number of layers
            input_initial_h: initial hidden state, tensor
            name: name of the variable scope to use
            reuse:whether to reuse already initialized variable

        Returns:
            h - all hidden states along T dimension,
                tf.Tensor with dimensionality [B x T x F]
            h_last - last hidden state, tf.Tensor with dimensionality [B x H]
        """
    with tf.variable_scope(name, reuse=reuse):

        if trainable_initial_states:
            init_h = tf.get_variable('init_h', [n_layers, 1, n_hidden])
            init_h = tf.tile(init_h, (1, tf.shape(units)[0], 1))
        else:
            init_h = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = input_initial_h or init_h

        with tf.variable_scope('cudnn_gru', reuse=reuse):
            def single_cell(): return tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell(n_hidden)
            cell = tf.nn.rnn_cell.MultiRNNCell([single_cell() for _ in range(n_layers)])

            units = tf.transpose(units, (1, 0, 2))

            h, h_last = tf.nn.dynamic_rnn(cell=cell, inputs=units, time_major=True,
                                          initial_state=tuple(tf.unstack(initial_h, axis=0)))
            h = tf.transpose(h, (1, 0, 2))

            h_last = h_last[-1]  # h_last is tuple: n_layers x batch_size x n_hidden

            # Extract last states if they are provided
            if seq_lengths is not None:
                indices = tf.stack([tf.range(tf.shape(h)[0]), seq_lengths-1], axis=1)
                h_last = tf.gather_nd(h, indices)

            return h, h_last
def cudnn_compatible_gru(units, n_hidden, n_layers=1, trainable_initial_states=False,
                         seq_lengths=None, input_initial_h=None, name='cudnn_gru', reuse=False):
    """ CuDNN Compatible GRU implementation.
        It should be used to load models saved with CudnnGRUCell to run on CPU.

        Args:
            units: tf.Tensor with dimensions [B x T x F], where
                B - batch size
                T - number of tokens
                F - features

            n_hidden: dimensionality of hidden state
            trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
            seq_lengths: tensor of sequence lengths with dimension [B]
            n_layers: number of layers
            input_initial_h: initial hidden state, tensor
            name: name of the variable scope to use
            reuse:whether to reuse already initialized variable

        Returns:
            h - all hidden states along T dimension,
                tf.Tensor with dimensionality [B x T x F]
            h_last - last hidden state, tf.Tensor with dimensionality [B x H]
        """
    with tf.variable_scope(name, reuse=reuse):

        if trainable_initial_states:
            init_h = tf.get_variable('init_h', [n_layers, 1, n_hidden])
            init_h = tf.tile(init_h, (1, tf.shape(units)[0], 1))
        else:
            init_h = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = input_initial_h or init_h

        with tf.variable_scope('cudnn_gru', reuse=reuse):
            def single_cell(): return tf.contrib.cudnn_rnn.CudnnCompatibleGRUCell(n_hidden)
            cell = tf.nn.rnn_cell.MultiRNNCell([single_cell() for _ in range(n_layers)])

            units = tf.transpose(units, (1, 0, 2))

            h, h_last = tf.nn.dynamic_rnn(cell=cell, inputs=units, time_major=True,
                                          initial_state=tuple(tf.unstack(initial_h, axis=0)))
            h = tf.transpose(h, (1, 0, 2))

            h_last = h_last[-1]  # h_last is tuple: n_layers x batch_size x n_hidden

            # Extract last states if they are provided
            if seq_lengths is not None:
                indices = tf.stack([tf.range(tf.shape(h)[0]), seq_lengths-1], axis=1)
                h_last = tf.gather_nd(h, indices)

            return h, h_last
CuDNN Compatible GRU implementation.
        It should be used to load models saved with CudnnGRUCell to run on CPU.

        Args:
            units: tf.Tensor with dimensions [B x T x F], where
                B - batch size
                T - number of tokens
                F - features

            n_hidden: dimensionality of hidden state
            trainable_initial_states: whether to create a special trainable variable
                to initialize the hidden states of the network or use just zeros
            seq_lengths: tensor of sequence lengths with dimension [B]
            n_layers: number of layers
            input_initial_h: initial hidden state, tensor
            name: name of the variable scope to use
            reuse:whether to reuse already initialized variable

        Returns:
            h - all hidden states along T dimension,
                tf.Tensor with dimensionality [B x T x F]
            h_last - last hidden state, tf.Tensor with dimensionality [B x H]