danche
danche

Reputation: 1815

How to share LSTM unit for 2 separate input in TensorFlow?

Assume I have 2 input q and a, how to make the 2 inputs share 1 LSTM cell? Now part of my code as belows

def lstmnets(self, sequence, seq_len):
    seq_embeds = self.embeds(sequence)

    # lstm_cell = tf.contrib.rnn.BasicLSTMCell(self.hidden_size)
    lstm_cell = tf.nn.rnn_cell.LSTMCell(self.hidden_size)
    init_state = lstm_cell.zero_state(self.batch_size, dtype=tf.float32)
    lstm_out, final_state = tf.nn.dynamic_rnn(lstm_cell, seq_embeds, initial_state=init_state, sequence_length=seq_len)
    return lstm_out

def inference(self, q, a, q_len, a_len):
    with tf.variable_scope('lstmnets') as scope:
        query_rep = self.lstmnets(q, q_len)
        scope.reuse_variables()
        title_rep = self.lstmnets(a, a_len)

But for this codes, My structure have 2 stacked LSTM as following figure. How can I just use one LSTM for both? In addition, how can I initial the LSTM weights and add them to histogram? so far, I find no related tutorials for this. Thanks.

enter image description here

Upvotes: 1

Views: 350

Answers (1)

Vijay Mariappan
Vijay Mariappan

Reputation: 17201

Your code seems to be fine as it uses scope.reuse_variable() to share the LSTM weights. The best way is to check is by printing the variables in the graph and verify whether the lstm_cell is declared only once. So in your inference function print the variable names:

def inference(self, q, a, q_len, a_len):
  with tf.variable_scope('lstmnets') as scope:
    query_rep = self.lstmnets(q, q_len)
    scope.reuse_variables()
    title_rep = self.lstmnets(a, a_len)
  for v in tf.global_variables():
    print(v.name)

Upvotes: 1

Related Questions