mediator.py

import tensorflow as tf
from tensorflow.contrib.cudnn_rnn.python.layers.cudnn_rnn import CudnnLSTM, CUDNN_RNN_BIDIRECTION
from tensorflow.contrib import layers
class Critic(object):
    def __call__(self, h):
        # sequence -> [b, l, v]
        _, l, v = h.get_shape().as_list()
        h = tf.reshape(h, [-1, l, 1, v])
        with tf.variable_scope("textmover", reuse=tf.AUTO_REUSE):
            h0 = layers.convolution2d(
                h, v, [4, 1], [2, 1],
                activation_fn=tf.nn.softplus
            )
            h1 = layers.convolution2d(
                h0, v, [4, 1], [1, 1],
                activation_fn=tf.nn.softplus
            )
            h2 = layers.convolution2d(
                h1, v, [4, 1], [2, 1],
                activation_fn=tf.nn.softplus
            )
            h = layers.flatten(h2)
            h = layers.fully_connected(
                h, 1, activation_fn=tf.identity
            )
            return h

class Mediator(object):
    def __init__(self, num_emb, batch_size, emb_dim, hidden_dim,
                 sequence_length, start_token,
                 learning_rate=1e-3, reward_gamma=0.95, name="mediator", dropout_rate=0.5, with_professor_forcing=False):
        self.num_emb = num_emb
        # self.batch_size = batch_size
        self.emb_dim = emb_dim
        self.hidden_dim = hidden_dim
        self.sequence_length = sequence_length
        self.learning_rate = tf.Variable(float(learning_rate), trainable=False)
        self.reward_gamma = reward_gamma
        self.g_params = []
        self.d_params = []
        self.temperature = 1.0
        self.name = name
        self.dropout_keep_rate = tf.Variable(float(1.0), trainable=False)
        self.dropout_on = self.dropout_keep_rate.assign(dropout_rate)
        self.dropout_off = self.dropout_keep_rate.assign(1.0)
        self.expected_reward = tf.Variable(tf.zeros([self.sequence_length]))

        self.x0 = tf.placeholder(tf.int32, shape=[None, self.sequence_length])
        self.x = self.x0
        self.x1 = tf.placeholder(tf.int32, shape=[None, self.sequence_length])
        input_x0 = tf.pad(self.x0, [[0, 0], [1, 0]])[:, 0:self.sequence_length]
        input_x1 = tf.pad(self.x1, [[0, 0], [1, 0]])[:, 0:self.sequence_length]
        output_x0 = tf.one_hot(
            self.x0, self.num_emb, on_value=1.0, off_value=0.0
        )
        output_x1 = tf.one_hot(
            self.x1, self.num_emb, on_value=1.0, off_value=0.0
        )
        with tf.variable_scope(name, reuse=tf.AUTO_REUSE):
            embedding = tf.get_variable(
                name="word_embeddings",
                initializer=tf.random_normal(shape=[self.num_emb, self.emb_dim], stddev=0.1)
            )
            Wo = tf.get_variable(
                name="Weight_output",
                initializer=tf.random_normal(shape=[self.hidden_dim, self.num_emb], stddev=0.1)
            )
            bo = tf.get_variable(
                name="bias_output",
                initializer=tf.random_normal(shape=[self.num_emb], stddev=0.1)
            )
            rnn = CudnnLSTM(
                num_layers=1,
                num_units=self.hidden_dim,
                kernel_initializer=tf.orthogonal_initializer()
            )
            def language_modeling(input_x):
                with tf.variable_scope("language_model", reuse=tf.AUTO_REUSE):
                    emb_x = tf.nn.embedding_lookup(
                        embedding, input_x
                    )
                    emb_x = tf.transpose(emb_x, [1, 0, 2])
                    h, _ = rnn(emb_x)
                    h = tf.transpose(h, [1, 0, 2])
                    h = tf.nn.dropout(h, self.dropout_keep_rate)
                    pred = tf.nn.log_softmax(
                        tf.reshape(h, [-1, self.hidden_dim]) @ Wo + bo,
                        axis=-1)
                    return h, tf.reshape(pred, [-1, self.sequence_length, self.num_emb])
            self.h0, self.log_predictions = language_modeling(input_x0)
            self.h1, self.log_predictions_ = language_modeling(input_x1)
            self.likelihood_loss = -tf.reduce_mean(
                tf.reduce_sum(
                    self.log_predictions * output_x0 +
                    self.log_predictions_ * output_x1, axis=-1)
            ) / 2.0
        self.m_opt = tf.train.AdamOptimizer(learning_rate, beta1=0.9, beta2=0.95)
        if with_professor_forcing:
            with tf.variable_scope("professor_forcing", reuse=tf.AUTO_REUSE):
                critic = Critic()
                myu = tf.random_uniform(shape=[tf.shape(self.x0)[0], self.sequence_length, 1],
                                        minval=0.0, maxval=1.0)
                hybrid = self.h0 * myu + self.h1 * (1.0 - myu)
                gp = tf.reduce_mean(tf.nn.relu(tf.norm(
                    tf.reshape(tf.gradients(critic(hybrid), [hybrid])[0], [tf.shape(self.x0)[0], -1]),
                    axis=-1) - 1.0) ** 2)
                self.d_loss = tf.reduce_mean(critic(self.h0) - critic(self.h1))
                self.d_opt = tf.train.AdamOptimizer(1e-4, beta1=0.5, beta2=0.9)
                self.d_params = [v for v in tf.trainable_variables() if "professor_forcing" in v.name]
                self.d_update = self.d_opt.minimize(self.d_loss + 5.0 * gp, var_list=self.d_params)
        self.m_params = [v for v in tf.trainable_variables() if name in v.name]
        if not with_professor_forcing:
            self.likelihood_updates = self.m_opt.minimize(self.likelihood_loss, var_list=self.m_params)
        else:
            self.likelihood_updates = self.m_opt.minimize(self.likelihood_loss - self.d_loss, var_list=self.m_params)

    def get_reward(self, sess, x):
        output = sess.run(self.log_predictions, feed_dict={self.x0: x})
        return output