tensorflow_test

tensorflow_test.py

import tensorflow as tf
import numpy as np
import time

from tensorflow.contrib.rnn import BasicLSTMCell

# https://github.com/tensorflow/tensorflow/issues/24828
try:
    from tensorflow.compat.v1 import ConfigProto
    from tensorflow.compat.v1 import InteractiveSession
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
except:
    pass

N = 1
L = 500
d = 10
k = 1
V = 1000

inputs = tf.placeholder('int64', [N, L])
emb_mat = tf.get_variable('emb_mat', shape=[V, d])
x = tf.nn.embedding_lookup(emb_mat, inputs)
lstm_cell = BasicLSTMCell(d)
a, (c, h) = tf.nn.dynamic_rnn(lstm_cell, x, dtype='float')
init_op = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init_op)

    inputs_val = np.zeros([N, L], dtype='int64')
    t0 = time.time()
    for _ in range(k):
        out = sess.run(h, feed_dict={inputs: inputs_val})
    t1 = time.time()
    print(t1 - t0)