deeppavlov/core/layers/tf_layers.py
def cudnn_compatible_lstm(units, n_hidden, n_layers=1, trainable_initial_states=None, seq_lengths=None, initial_h=None,
                          initial_c=None, name='cudnn_lstm', reuse=False):
    """ CuDNN Compatible LSTM implementation.
        It should be used to load models saved with CudnnLSTMCell 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
            n_layers: number of layers
            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]
            initial_h: optional initial hidden state, masks trainable_initial_states
                if provided
            initial_c: optional initial cell state, masks trainable_initial_states
                if provided
            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]
                where H - number of hidden units
            c_last - last cell state, tf.Tensor with dimensionality [B x H]
                where H - number of hidden units
        """

    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))
            init_c = tf.get_variable('init_c', [n_layers, 1, n_hidden])
            init_c = tf.tile(init_c, (1, tf.shape(units)[0], 1))
        else:
            init_h = init_c = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = initial_h or init_h
        initial_c = initial_c or init_c

        with tf.variable_scope('cudnn_lstm', reuse=reuse):
            def single_cell(): return tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell(n_hidden)

            cell = tf.nn.rnn_cell.MultiRNNCell([single_cell() for _ in range(n_layers)])

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

            init = tuple([tf.nn.rnn_cell.LSTMStateTuple(ic, ih) for ih, ic in
                          zip(tf.unstack(initial_h, axis=0), tf.unstack(initial_c, axis=0))])

            h, state = tf.nn.dynamic_rnn(cell=cell, inputs=units, time_major=True, initial_state=init)

            h = tf.transpose(h, (1, 0, 2))
            h_last = state[-1].h
            c_last = state[-1].c

            # 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, c_last)
def cudnn_compatible_lstm(units, n_hidden, n_layers=1, trainable_initial_states=None, seq_lengths=None, initial_h=None,
                          initial_c=None, name='cudnn_lstm', reuse=False):
    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))
            init_c = tf.get_variable('init_c', [n_layers, 1, n_hidden])
            init_c = tf.tile(init_c, (1, tf.shape(units)[0], 1))
        else:
            init_h = init_c = tf.zeros([n_layers, tf.shape(units)[0], n_hidden])

        initial_h = initial_h or init_h
        initial_c = initial_c or init_c

        with tf.variable_scope('cudnn_lstm', reuse=reuse):
            def single_cell(): return tf.contrib.cudnn_rnn.CudnnCompatibleLSTMCell(n_hidden)

            cell = tf.nn.rnn_cell.MultiRNNCell([single_cell() for _ in range(n_layers)])

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

            init = tuple([tf.nn.rnn_cell.LSTMStateTuple(ic, ih) for ih, ic in
                          zip(tf.unstack(initial_h, axis=0), tf.unstack(initial_c, axis=0))])

            h, state = tf.nn.dynamic_rnn(cell=cell, inputs=units, time_major=True, initial_state=init)

            h = tf.transpose(h, (1, 0, 2))
            h_last = state[-1].h
            c_last = state[-1].c

            # 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, c_last)
def cudnn_compatible_lstm ( units , n_hidden , n_layers = 1 , trainable_initial_states = None , seq_lengths = None , initial_h = None , initial_c = None , name = 'cudnn_lstm' , reuse = False ) : 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 ) ) init_c = tf . get_variable ( 'init_c' , [ n_layers , 1 , n_hidden ] ) init_c = tf . tile ( init_c , ( 1 , tf . shape ( units ) [ 0 ] , 1 ) ) else : init_h = init_c = tf . zeros ( [ n_layers , tf . shape ( units ) [ 0 ] , n_hidden ] ) initial_h = initial_h or init_h initial_c = initial_c or init_c with tf . variable_scope ( 'cudnn_lstm' , reuse = reuse ) : def single_cell ( ) : return tf . contrib . cudnn_rnn . CudnnCompatibleLSTMCell ( n_hidden ) cell = tf . nn . rnn_cell . MultiRNNCell ( [ single_cell ( ) for _ in range ( n_layers ) ] ) units = tf . transpose ( units , ( 1 , 0 , 2 ) ) init = tuple ( [ tf . nn . rnn_cell . LSTMStateTuple ( ic , ih ) for ih , ic in zip ( tf . unstack ( initial_h , axis = 0 ) , tf . unstack ( initial_c , axis = 0 ) ) ] ) h , state = tf . nn . dynamic_rnn ( cell = cell , inputs = units , time_major = True , initial_state = init ) h = tf . transpose ( h , ( 1 , 0 , 2 ) ) h_last = state [ - 1 ] . h c_last = state [ - 1 ] . c # 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 , c_last )
CuDNN Compatible LSTM implementation.
        It should be used to load models saved with CudnnLSTMCell 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
            n_layers: number of layers
            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]
            initial_h: optional initial hidden state, masks trainable_initial_states
                if provided
            initial_c: optional initial cell state, masks trainable_initial_states
                if provided
            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]
                where H - number of hidden units
            c_last - last cell state, tf.Tensor with dimensionality [B x H]
                where H - number of hidden units
CuDNN Compatible LSTM implementation . It should be used to load models saved with CudnnLSTMCell to run on CPU .