#!/usr/bin/env python from six.moves import xrange # pylint: disable=redefined-builtin import math import time import numpy as np from datetime import datetime import reader import tensorflow as tf from tensorflow.python.ops import rnn FLAGS = tf.app.flags.FLAGS tf.app.flags.DEFINE_integer('batch_size', 128, """Batch size.""") tf.app.flags.DEFINE_integer('num_batches', 100, """Number of batches to run.""") tf.app.flags.DEFINE_integer('num_layers', 1, """Number of batches to run.""") tf.app.flags.DEFINE_integer('max_len', 100, """Number of batches to run.""") tf.app.flags.DEFINE_boolean('forward_only', False, """Only run the forward pass.""") tf.app.flags.DEFINE_boolean('forward_backward_only', False, """Only run the forward-forward pass.""") tf.app.flags.DEFINE_integer('hidden_size', 128, """Number of batches to run.""") tf.app.flags.DEFINE_integer('emb_size', 128, """Number of batches to run.""") tf.app.flags.DEFINE_boolean('log_device_placement', False, """Whether to log device placement.""") VOCAB_SIZE=30000 NUM_CLASS=2 def get_feed_dict(x_data, y_data=None): feed_dict = {} if y_data is not None: feed_dict[y_input] = y_data for i in xrange(x_data.shape[0]): feed_dict[x_input[i]] = x_data[i, :, :] return feed_dict def get_incoming_shape(incoming): """ Returns the incoming data shape """ if isinstance(incoming, tf.Tensor): return incoming.get_shape().as_list() elif type(incoming) in [np.array, list, tuple]: return np.shape(incoming) else: raise Exception("Invalid incoming layer.") # Note input * W is done in LSTMCell, # which is different from PaddlePaddle def single_lstm(name, incoming, n_units, use_peepholes=True, return_seq=False, return_state=False): with tf.name_scope(name) as scope: cell = tf.nn.rnn_cell.LSTMCell(n_units, use_peepholes=use_peepholes) output, _cell_state = rnn.rnn(cell, incoming, dtype=tf.float32) out = output if return_seq else output[-1] return (out, _cell_state) if return_state else out def lstm(name, incoming, n_units, use_peepholes=True, return_seq=False, return_state=False, num_layers=1): with tf.name_scope(name) as scope: lstm_cell = tf.nn.rnn_cell.LSTMCell(n_units, use_peepholes=use_peepholes) cell = tf.nn.rnn_cell.MultiRNNCell([lstm_cell] * num_layers) initial_state = cell.zero_state(FLAGS.batch_size, dtype=tf.float32) if not isinstance(incoming, list): # if the input is embeding, the Tensor shape : [None, time_step, emb_size] incoming = [tf.squeeze(input_, [1]) for input_ in tf.split(1, FLAGS.max_len, incoming)] outputs, state = tf.nn.rnn(cell, incoming, initial_state=initial_state, dtype=tf.float32) out = outputs if return_seq else outputs[-1] return (out, _cell_state) if return_state else out def embedding(name, incoming, vocab_size, emb_size): with tf.name_scope(name) as scope: #with tf.device("/cpu:0"): embedding = tf.get_variable( name+'_emb', [vocab_size, emb_size], dtype=tf.float32) out = tf.nn.embedding_lookup(embedding, incoming) return out def fc(name, inpOp, nIn, nOut, act=True): with tf.name_scope(name) as scope: kernel = tf.get_variable(name + '_w', [nIn, nOut], initializer=tf.truncated_normal_initializer(stddev=0.01, dtype=tf.float32), dtype=tf.float32) biases = tf.get_variable(name + '_b', [nOut], initializer=tf.constant_initializer(value=0.0, dtype=tf.float32), dtype=tf.float32,trainable=True) net = tf.nn.relu_layer(inpOp, kernel, biases, name=name) if act else \ tf.matmul(inpOp, kernel) + biases return net def inference(seq): net = embedding('emb', seq, VOCAB_SIZE, FLAGS.emb_size) print "emb:", get_incoming_shape(net) net = lstm('lstm', net, FLAGS.hidden_size, num_layers=FLAGS.num_layers) print "lstm:", get_incoming_shape(net) net = fc('fc1', net, FLAGS.hidden_size, 2) return net def loss(logits, labels): # one label index for one sample labels = tf.cast(labels, tf.float32) cross_entropy = tf.nn.softmax_cross_entropy_with_logits( logits, labels, name='cross_entropy_per_example') cross_entropy_mean = tf.reduce_mean(cross_entropy, name='cross_entropy') tf.add_to_collection('losses', cross_entropy_mean) return tf.add_n(tf.get_collection('losses'), name='total_loss') def time_tensorflow_run(session, target, x_input, y_input, info_string): num_steps_burn_in = 50 total_duration = 0.0 total_duration_squared = 0.0 if not isinstance(target, list): target = [target] target_op = tf.group(*target) train_dataset = reader.create_datasets("imdb.pkl", VOCAB_SIZE) for i in xrange(FLAGS.num_batches + num_steps_burn_in): start_time = time.time() data, label = train_dataset.next_batch(FLAGS.batch_size) _ = session.run(target_op, feed_dict={x_input:data, y_input:label}) duration = time.time() - start_time if i > num_steps_burn_in: if not i % 10: print ('%s: step %d, duration = %.3f' % (datetime.now(), i - num_steps_burn_in, duration)) total_duration += duration total_duration_squared += duration * duration mn = total_duration / FLAGS.num_batches vr = total_duration_squared / FLAGS.num_batches - mn * mn sd = math.sqrt(vr) print ('%s: %s across %d steps, %.3f +/- %.3f sec / batch' % (datetime.now(), info_string, FLAGS.num_batches, mn, sd)) def run_benchmark(): with tf.Graph().as_default(): global_step=0 with tf.device('/cpu:0'): global_step = tf.Variable(0, trainable=False) with tf.device('/gpu:0'): #x_input = tf.placeholder(tf.int32, [None, FLAGS.max_len], name="x_input") #y_input = tf.placeholder(tf.int32, [None, NUM_CLASS], name="y_input") x_input = tf.placeholder(tf.int32, [FLAGS.batch_size, FLAGS.max_len], name="x_input") y_input = tf.placeholder(tf.int32, [FLAGS.batch_size, NUM_CLASS], name="y_input") # Generate some dummy sequnce. last_layer = inference(x_input) objective = loss(last_layer, y_input) opt = tf.train.AdamOptimizer(0.001) grads = opt.compute_gradients(objective) apply_gradient_op = opt.apply_gradients(grads, global_step=global_step) init = tf.initialize_all_variables() sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True, log_device_placement=FLAGS.log_device_placement)) sess.run(init) run_forward = True run_forward_backward = True if FLAGS.forward_only and FLAGS.forward_backward_only: raise ValueError("Cannot specify --forward_only and " "--forward_backward_only at the same time.") if FLAGS.forward_only: run_forward_backward = False elif FLAGS.forward_backward_only: run_forward = False if run_forward: time_tensorflow_run(sess, last_layer, x_input, y_input, "Forward") if run_forward_backward: with tf.control_dependencies([apply_gradient_op]): train_op = tf.no_op(name='train') time_tensorflow_run(sess, [train_op, objective], x_input, y_input, "Forward-backward") def main(_): run_benchmark() if __name__ == '__main__': tf.app.run()