理工男的文科梦 —— LSTM深度学习写春联

couplet_file ="couplet.txt"#对联couplets = []with open(couplet_file,'r') as f:    for line in f:        try:
            content = line.replace(' ','')            if '_' in content or '(' in content or '（' in content or '《' in content or '[' in content:                continue
            if len(content) < 5*3 or len(content) > 79*3:                continue
            content = '[' + content + ']'
           # print chardet.detect(content)
            content = content.decode('utf-8')
            couplets.append(content)        except Exception as e:            pass# 按字数排序couplets = sorted(couplets,key=lambda line: len(line))
print('对联总数: %d'%(len(couplets)))# 统计每个字出现次数all_words = []for couplet in couplets:
    all_words += [word for word in couplet]

counter = collections.Counter(all_words)
count_pairs = sorted(counter.items(), key=lambda x: -x[1])
words, _ = zip(*count_pairs)
words = words[:len(words)] + (' ',)# 每个字映射为一个数字IDword_num_map = dict(zip(words, range(len(words))))

to_num = lambda word: word_num_map.get(word, len(words))
couplets_vector = [ list(map(to_num, couplet)) for couplet in couplets]# 每次取64首对联进行训练, 此参数可以调整batch_size = 64n_chunk = len(couplets_vector) // batch_size
x_batches = []
y_batches = []for i in range(n_chunk):
    start_index = i * batch_size#起始位置
    end_index = start_index + batch_size#结束位置

    batches = couplets_vector[start_index:end_index]
    length = max(map(len,batches))#每个batches中句子的最大长度
    xdata = np.full((batch_size,length), word_num_map[' '], np.int32)    for row in range(batch_size):
        xdata[row,:len(batches[row])] = batches[row]
    ydata = np.copy(xdata)
    ydata[:,:-1] = xdata[:,1:]
    x_batches.append(xdata)
    y_batches.append(ydata

def neural_network(rnn_size=128, num_layers=2):
    cell = tf.nn.rnn_cell.BasicLSTMCell(rnn_size, state_is_tuple=True)
    cell = tf.nn.rnn_cell.MultiRNNCell([cell] * num_layers, state_is_tuple=True)

    initial_state = cell.zero_state(batch_size, tf.float32)    with tf.variable_scope('rnnlm'):
        softmax_w = tf.get_variable("softmax_w", [rnn_size, len(words)+1])
        softmax_b = tf.get_variable("softmax_b", [len(words)+1])        with tf.device("/cpu:0"):
            embedding = tf.get_variable("embedding", [len(words)+1, rnn_size])
            inputs = tf.nn.embedding_lookup(embedding, input_data)

    outputs, last_state = tf.nn.dynamic_rnn(cell, inputs, initial_state=initial_state, scope='rnnlm')
    output = tf.reshape(outputs,[-1, rnn_size])

    logits = tf.matmul(output, softmax_w) + softmax_b
    probs = tf.nn.softmax(logits)    return logits, last_state, probs, cell, initial_state

def train_neural_network():
    logits, last_state, _, _, _ = neural_network()
    targets = tf.reshape(output_targets, [-1])
    loss = tf.contrib.legacy_seq2seq.sequence_loss_by_example([logits], [targets], [tf.ones_like(targets, dtype=tf.float32)], len(words))
    cost = tf.reduce_mean(loss)
    learning_rate = tf.Variable(0.0, trainable=False)
    tvars = tf.trainable_variables()
    grads, _ = tf.clip_by_global_norm(tf.gradients(cost, tvars), 5)
    optimizer = tf.train.AdamOptimizer(learning_rate)
    train_op = optimizer.apply_gradients(zip(grads, tvars))    with tf.Session() as sess:
        sess.run(tf.initialize_all_variables())
        saver = tf.train.Saver(tf.all_variables())        for epoch in range(100):
            sess.run(tf.assign(learning_rate, 0.01 * (0.97 ** epoch)))
            n = 0
            for batche in range(n_chunk):
                train_loss, _ , _ = sess.run([cost, last_state, train_op], feed_dict={input_data: x_batches[n], output_targets: y_batches[n]})
                n += 1
                print(epoch, batche, train_loss)            if epoch % 7 == 0:
                saver.save(sess, './couplet.module', global_step=epoch)