|
| 1 | +import codecs |
| 2 | +import numpy as np |
| 3 | +import logging |
| 4 | +import tensorflow as tf |
| 5 | +from tensorflow.contrib import rnn |
| 6 | +import sys |
| 7 | + |
| 8 | +data_file = "data/rnn-train-data.txt" |
| 9 | +rnn_layers = 2 |
| 10 | +embedding_size = 128 |
| 11 | +hidden_size = 128 |
| 12 | +input_dropout = 0.2 |
| 13 | +learning_rate = 0.01 |
| 14 | +max_grad_norm = 5 |
| 15 | +num_epochs = 500001 |
| 16 | +batch_size = 20 |
| 17 | +seq_length = 10 |
| 18 | + |
| 19 | + |
| 20 | +def main(): |
| 21 | + logging.basicConfig(stream=sys.stdout, |
| 22 | + format='%(asctime)s %(levelname)s:%(message)s', |
| 23 | + level=logging.INFO, |
| 24 | + datefmt='%I:%M:%S') |
| 25 | + with codecs.open(data_file, 'r') as f: |
| 26 | + text = f.read() |
| 27 | + train_text = text |
| 28 | + vocab_index_dict, index_vocab_dict, vocab_size = create_vocab(text) |
| 29 | + |
| 30 | + train_batches = BatchGenerator(train_text, batch_size, seq_length, vocab_size, vocab_index_dict) |
| 31 | + graph = tf.Graph() |
| 32 | + with graph.as_default(): |
| 33 | + input_data = tf.placeholder(tf.int64, [batch_size, seq_length], name='inputs') |
| 34 | + input_targets = tf.placeholder(tf.int64, [batch_size, seq_length], name='targets') |
| 35 | + tf_learning_rate = tf.constant(learning_rate) |
| 36 | + |
| 37 | + embedding = tf.get_variable('embedding', [vocab_size, embedding_size]) |
| 38 | + inputs = tf.nn.embedding_lookup(embedding, input_data) |
| 39 | + sliced_inputs = [tf.squeeze(input_, [1]) for input_ in |
| 40 | + tf.split(axis=1, num_or_size_splits=seq_length, value=inputs)] |
| 41 | + |
| 42 | + weights = tf.Variable(tf.random_normal([2 * hidden_size, vocab_size])) |
| 43 | + biases = tf.Variable(tf.random_normal([vocab_size])) |
| 44 | + |
| 45 | + lstm_fw_cell = rnn.BasicLSTMCell(hidden_size, forget_bias=1.0) |
| 46 | + lstm_bw_cell = rnn.BasicLSTMCell(hidden_size, forget_bias=1.0) |
| 47 | + outputs, _, _ = rnn.static_bidirectional_rnn(lstm_fw_cell, lstm_bw_cell, sliced_inputs, dtype=tf.float32) |
| 48 | + |
| 49 | + flat_outputs = tf.reshape(tf.concat(axis=1, values=outputs), [-1, 2 * hidden_size]) |
| 50 | + flat_targets = tf.reshape(tf.concat(axis=1, values=input_targets), [-1]) |
| 51 | + logits = tf.matmul(flat_outputs, weights) + biases |
| 52 | + loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=flat_targets) |
| 53 | + mean_loss = tf.reduce_mean(loss) |
| 54 | + |
| 55 | + tvars = tf.trainable_variables() |
| 56 | + grads, _ = tf.clip_by_global_norm(tf.gradients(mean_loss, tvars), max_grad_norm) |
| 57 | + optimizer = tf.train.AdamOptimizer(tf_learning_rate) |
| 58 | + train_op = optimizer.apply_gradients(zip(grads, tvars)) |
| 59 | + |
| 60 | + prediction = tf.nn.softmax(logits) |
| 61 | + correct_pred = tf.equal(tf.argmax(prediction, 1), flat_targets) |
| 62 | + accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32)) |
| 63 | + |
| 64 | + with tf.Session(graph=graph) as session: |
| 65 | + tf.global_variables_initializer().run() |
| 66 | + for i in range(num_epochs): |
| 67 | + data = train_batches.next() |
| 68 | + inputs = np.array(data[:-1]).transpose() |
| 69 | + targets = np.array(data[1:]).transpose() |
| 70 | + ops = [mean_loss, train_op, tf_learning_rate, accuracy] |
| 71 | + feed_dict = {input_data: inputs, input_targets: targets} |
| 72 | + average_loss, __, lr, acc = session.run(ops, feed_dict) |
| 73 | + if i % 100 == 0: |
| 74 | + logging.info("average loss: %.5f,accuracy: %.3f", average_loss, acc) |
| 75 | + |
| 76 | + |
| 77 | +def create_vocab(text): |
| 78 | + unique_chars = list(set(text)) |
| 79 | + print(unique_chars) |
| 80 | + vocab_size = len(unique_chars) |
| 81 | + vocab_index_dict = {} |
| 82 | + index_vocab_dict = {} |
| 83 | + for i, char in enumerate(unique_chars): |
| 84 | + vocab_index_dict[char] = i |
| 85 | + index_vocab_dict[i] = char |
| 86 | + return vocab_index_dict, index_vocab_dict, vocab_size |
| 87 | + |
| 88 | + |
| 89 | +class BatchGenerator(object): |
| 90 | + def __init__(self, text, batch_size, seq_length, vocab_size, vocab_index_dict): |
| 91 | + self._text = text |
| 92 | + self._text_size = len(text) |
| 93 | + self._batch_size = batch_size |
| 94 | + self.vocab_size = vocab_size |
| 95 | + self.seq_length = seq_length |
| 96 | + self.vocab_index_dict = vocab_index_dict |
| 97 | + |
| 98 | + segment = self._text_size // batch_size |
| 99 | + |
| 100 | + self._cursor = [offset * segment for offset in range(batch_size)] |
| 101 | + self._last_batch = self._next_batch() |
| 102 | + |
| 103 | + def _next_batch(self): |
| 104 | + batch = np.zeros(shape=(self._batch_size), dtype=np.float) |
| 105 | + for b in range(self._batch_size): |
| 106 | + batch[b] = self.vocab_index_dict[self._text[self._cursor[b]]] |
| 107 | + self._cursor[b] = (self._cursor[b] + 1) % self._text_size |
| 108 | + return batch |
| 109 | + |
| 110 | + def next(self): |
| 111 | + batches = [self._last_batch] |
| 112 | + for step in range(self.seq_length): |
| 113 | + batches.append(self._next_batch()) |
| 114 | + self._last_batch = batches[-1] |
| 115 | + return batches |
| 116 | + |
| 117 | + |
| 118 | +if __name__ == '__main__': |
| 119 | + main() |
0 commit comments