Huge second-derivative graph

huge_graph.py

"""
Generate some second-derivatives that make TF cry.

On my machine, this takes several minutes to run and uses
about 2GB of memory.
"""

import time

import tensorflow as tf

def build_graph():
    """
    Build a graph that TF hates.
    """
    print('Building RNN...')
    inputs = [tf.Variable(tf.random_normal((10, 5))) for _ in range(200)]
    rnn_out, _, _ = tf.nn.static_bidirectional_rnn(tf.contrib.rnn.LSTMCell(128),
                                                   tf.contrib.rnn.LSTMCell(128),
                                                   inputs,
                                                   dtype=tf.float32)
    rnn_out = tf.reduce_mean(tf.concat(rnn_out, 0))

    print('Building meta-loss...')
    targets = [tf.placeholder(tf.float32, (10, 5)) for _ in inputs]
    grads = tf.gradients(rnn_out, inputs)
    meta_losses = [tf.reduce_mean(tf.square(target - grad))
                   for target, grad in zip(targets, grads)]
    meta_loss = tf.reduce_sum(tf.stack(meta_losses))

    print('Building gradient of meta-loss...')
    tf.train.AdamOptimizer().minimize(meta_loss)

    print('Done building graph.')
    while True:
        time.sleep(1)


if __name__ == '__main__':
    build_graph()