Внедрить настроенный LSTM в Tensorflow - PullRequest
Новая
       52

Внедрить настроенный LSTM в Tensorflow

0 голосов
/ 27 сентября 2019

Я пытаюсь реализовать LSTM с нормализацией и маскированием слоев в Tensorflow.Вот мой код, примерно следующий за реализацией OpenAI с некоторыми изменениями. 

def lstm_norm(xs, units, masks):
    """lstm with masks and layer normalization

    Arguments:
        xs      3d Tensor   --  input data of shape [n_batch, n_seq, dim]
        units   int         --  size of hidden/cell state
        masks   2d Tensor   --  masks, must match the first 2 dimensions of xs

    Returns:
        ys      3d Tensor   -- output date of shape [n_batch, n_seq, units]
        initial_state, final_state
    """
    def ln(x, gamma, beta, eps=1e-5, axes=[1]):
        mean, var = tf.nn.moments(x, axes=axes, keep_dims=True)
        x = (x - mean) / tf.sqrt(var + eps)
        x = gamma * x + beta
        return x

    n_batch, n_steps, x_dim = xs.shape.as_list()

    xw_shape = [x_dim, 4*units]
    xb_shape = [4*units]
    hw_shape = [units, 4*units]
    hb_shape = [4*units]

    with tf.variable_scope('lstm_norm'):
        x_w = tf.get_variable('x_w', shape=xw_shape)
        x_g = tf.get_variable('x_g', [4*units], initializer=tf.constant_initializer(1.0))
        x_b = tf.get_variable('x_b', shape=xb_shape, 
                              initializer=tf.zeros_initializer())

        h_w = tf.get_variable('h_w', shape=hw_shape)
        h_g = tf.get_variable('h_g', [4*units], initializer=tf.constant_initializer(1.0))
        h_b = tf.get_variable('h_b', shape=hb_shape, 
                              initializer=tf.zeros_initializer())

        b = tf.get_variable('b', [4*units], initializer=tf.zeros_initializer())

        c_g = tf.get_variable('c_g', [units], initializer=tf.constant_initializer(1.0))
        c_b = tf.get_variable('c_b', [units], initializer=tf.zeros_initializer())

        initial_state = tf.zeros([n_batch, 2*units], name='initial_state')
        h, c = tf.split(value=initial_state, num_or_size_splits=2, axis=1)

        # unstack xs and masks so that we can iterate through next
        xs = tf.unstack(xs, n_steps, axis=1)
        if masks is not None:
            masks = tf.unstack(masks, n_steps, axis=1)

        ys = []
        for x, m in zip(xs, masks):
            m = m[..., None]
            c *= m
            h *= m
            # implementation of lstm with layer norm starts here
            z = ln(tf.matmul(x, x_w), x_g, x_b) + ln(tf.matmul(h, h_w), h_g, h_b)
            f, i, o, u = tf.split(value=z, num_or_size_splits=4, axis=1)
            f = tf.nn.sigmoid(f)
            i = tf.nn.sigmoid(i)
            o = tf.nn.sigmoid(o)
            u = tf.tanh(u)
            c = f * c + i * u
            h = o * tf.tanh(ln(c, c_g, c_b))
            ys.append(h)

        final_state = (h, c)
        ys = tf.stack(ys, 1)

    return ys, (initial_state, final_state)

tf.reset_default_graph()
xs = tf.placeholder(tf.float32, [64, 1000, 256])
masks = tf.placeholder(tf.float32, [64, 1000])
start = time.time()
# takes 36s to construct the graph
ys, _ =  lstm_norm(xs, 126, masks)
print(time.time() - start)

Этот код занимает значительно больше времени для построения графика, чем официальная реализация Tensorflow в следующем (36 с против 180 мс намой ноутбук без CUDA GPU).

# official implementation of LSTM, taking 180ms to construct the graph
tf.reset_default_graph()
xs = tf.placeholder(tf.float32, [64, 1000, 256])
lstm_keras = tf.keras.layers.LSTM(126, return_sequences=True, return_state=True)
ys_keras, _, _ = lstm_keras(xs)

Я также проверяю их время выполнения с помощью следующего кода 

# use my implementation of LSTM, taking 24s
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    start = time.time()
    y = sess.run(ys, feed_dict={xs: np.random.rand(64, 1000, 256), masks: np.random.randint(2, size=(64, 1000))})
    print(time.time() - start)
    print(y.shape)

# use the official implementation, using 0.8s
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    start = time.time()
    y = sess.run(ys_keras, feed_dict={xs: np.random.rand(64, 1000, 256)})
    print(time.time() - start)
    print(y.shape)

Кроме того, моя реализация имеет существенный недостаток, заключающийся в том, что она должна знать длину последовательности перед выполнением.Это ограничивает его использование в случаях, когда мы можем передавать последовательности различной длины - это можно обойти, повторно используя переменные, что, однако, усложнит код.Есть ли другой способ построить LSTM, который бы соответствовал моим потребностям?Более конкретно, я хочу реализовать LSTM со следующими свойствами

  1. Может применяться нормализация слоя
  2. Может применяться маскирование
  3. Может использоваться с различной длинойпоследовательности

Любой совет будет оценен.Заранее спасибо.

...