yusuf
yusuf

Reputation: 3781

(scan) function in theano and tensorflow

I have the following function in theano:

def forward_prop_step(x_t, s_t1_prev, s_t2_prev):
        # This is how we calculated the hidden state in a simple RNN. No longer!
        # s_t = T.tanh(U[:,x_t] + W.dot(s_t1_prev))

        # Word embedding layer
        x_e = E[:,x_t]

        # GRU Layer 1
        z_t1 = T.nnet.hard_sigmoid(U[0].dot(x_e) + W[0].dot(s_t1_prev) + b[0])
        r_t1 = T.nnet.hard_sigmoid(U[1].dot(x_e) + W[1].dot(s_t1_prev) + b[1])
        c_t1 = T.tanh(U[2].dot(x_e) + W[2].dot(s_t1_prev * r_t1) + b[2])
        s_t1 = (T.ones_like(z_t1) - z_t1) * c_t1 + z_t1 * s_t1_prev

        # GRU Layer 2
        z_t2 = T.nnet.hard_sigmoid(U[3].dot(s_t1) + W[3].dot(s_t2_prev) + b[3])
        r_t2 = T.nnet.hard_sigmoid(U[4].dot(s_t1) + W[4].dot(s_t2_prev) + b[4])
        c_t2 = T.tanh(U[5].dot(s_t1) + W[5].dot(s_t2_prev * r_t2) + b[5])
        s_t2 = (T.ones_like(z_t2) - z_t2) * c_t2 + z_t2 * s_t2_prev

        # Final output calculation
        # Theano's softmax returns a matrix with one row, we only need the row
        o_t = T.nnet.softmax(V.dot(s_t2) + c)[0]

        return [o_t, s_t1, s_t2]

And I call this function using scan:

[o, s, s2], updates = theano.scan(
            forward_prop_step,
            sequences=x,
            truncate_gradient=self.bptt_truncate,
            outputs_info=[None, 
                          dict(initial=T.zeros(self.hidden_dim)),
dict(initial=T.zeros(self.hidden_dim))])

I have tried to rewrite the same function in tensorflow:

def forward_prop_step(x_t, s_t1_prev, s_t2_prev):
     # Word embedding layer
     x_e = E[:, x_t]

     # GRU Layer 1
     z_t1 = tf.sigmoid(tf.reduce_sum(U[0] * x_e, axis=1) + tf.reduce_sum(W[0] * s_t1_prev, axis=1) + b[0])
     r_t1 = tf.sigmoid(tf.reduce_sum(U[1] * x_e, axis=1) + tf.reduce_sum(W[1] * s_t1_prev, axis=1) + b[1])
     c_t1 = tf.tanh(tf.reduce_sum(U[2] * x_e, axis=1) + tf.reduce_sum(W[2] * (s_t1_prev * r_t1), axis=1) + b[2])
     s_t1 = (tf.ones_like(z_t1) - z_t1) * c_t1 + z_t1 * s_t1_prev

     # GRU Layer 2
     z_t2 = tf.sigmoid(tf.reduce_sum(U[3] * s_t1, axis=1) + tf.reduce_sum(W[3] * s_t2_prev, axis=1) + b[3])
     r_t2 = tf.sigmoid(tf.reduce_sum(U[4] * s_t1, axis=1) + tf.reduce_sum(W[4] * s_t2_prev) + b[1])
     c_t2 = tf.tanh(tf.reduce_sum(U[5] * s_t1, axis=1) + tf.reduce_sum(W[5] * (s_t2_prev * r_t2), axis=1) + b[5])
     s_t2 = (tf.ones_like(z_t2) - z_t2) * c_t2 + z_t2 * s_t2_prev

     # Final output calculation
     o_t = tf.softmax(tf.reduce_sum(V * s_t2, axis=1) + c)[0]

     return [o_t, s_t1, s_t2]

And I have called this function using scan:

s = tf.zeros([self.hidden_dim, 0])
s2 = tf.zeros([self.hidden_dim, 0])

[o, s, s2] = tf.scan(
     fn=forward_prop_step,
     elems=[x, s, s2])

Instead of using initializer, I have initialized s and s2 variables before scan. When I run my code in tensorflow, I get the following error:

TypeError: forward_prop_step() takes exactly 3 arguments (2 given)

I am sure that the only problem is not the bug above. How can I rewrite scan function in tensorflow by getting reference the theano code?

Upvotes: 1

Views: 957

Answers (1)

Peter Hawkins
Peter Hawkins

Reputation: 3211

If you want to pass more than one element to tf.scan(), you need to wrap them up in a list or a tuple. Here's an example of how to do it:

def f(x, ys):
  (y1, y2) = ys
  return x + y1 * y2

a = tf.constant([1, 2, 3, 4, 5])
b = tf.constant([2, 3, 2, 2, 1])
c = tf.scan(f, (a, b), initializer=0)
with tf.Session() as sess:
  print(sess.run(c))

which prints:

[ 2  8 14 22 27]

I hope that helps!

Upvotes: 1

Related Questions