Occupies all GPU memory with the following tensorflow code?

Question

I was experimenting with word2vec with the code from https://github.com/chiphuyen/stanford-tensorflow-tutorials/blob/master/examples/04_word2vec_no_frills.py

However, it easily uses up all my GPU memories, any idea why?

with tf.name_scope('data'):
    center_words = tf.placeholder(tf.int32, shape=[BATCH_SIZE], name='center_words')
    target_words = tf.placeholder(tf.int32, shape=[BATCH_SIZE, 1], name='target_words')

with tf.name_scope("embedding_matrix"):
    embed_matrix = tf.Variable(tf.random_uniform([VOCAB_SIZE, EMBED_SIZE], -1.0, 1.0), name="embed_matrix")

with tf.name_scope("loss"):
    embed = tf.nn.embedding_lookup(embed_matrix, center_words, name="embed")

    nce_weight = tf.Variable(tf.truncated_normal([VOCAB_SIZE, EMBED_SIZE], stddev=1.0/(EMBED_SIZE ** 0.5)), name="nce_weight")
    nce_bias = tf.Variable(tf.zeros([VOCAB_SIZE]), name="nce_bias")


    loss = tf.reduce_mean(tf.nn.nce_loss(weights=nce_weight, biases=nce_bias, labels=target_words, inputs=embed, num_sampled=NUM_SAMPLED, num_classes=VOCAB_SIZE), name="loss")


optimizer = tf.train.GradientDescentOptimizer(LEARNING_RATE).minimize(loss)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    total_loss = 0.0 # we use this to calculate the average loss in the last SKIP_STEP steps
    writer = tf.summary.FileWriter('./graphs/no_frills/', sess.graph)
    for index in range(NUM_TRAIN_STEPS):
        centers, targets = next(batch_gen)
        loss_batch, _ = sess.run([loss, optimizer], feed_dict={center_words:centers, target_words:targets})

        total_loss += loss_batch
        if (index + 1) % SKIP_STEP == 0:
            print('Average loss at step {}: {:5.1f}'.format(index, total_loss / SKIP_STEP))
            total_loss = 0.0
    writer.close()

I think this is standard behavior for TF. See [here](https://stackoverflow.com/questions/34199233/how-to-prevent-tensorflow-from-allocating-the-totality-of-a-gpu-memory). If you google for example "tensorflow allocates all gpu memory" you will find other descriptions of it. — Robert Crovella, Sep 06 '17 at 23:57

score 5 · Answer 1 · answered Sep 07 '17 at 01:00

This is default tensorflow behaviour. If you want to limit GPU memory allocation to only what is needed, specify this in the session config.

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)

Alternatively you can specify a maximum fraction of GPU memory to use:

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.5
session = tf.Session(config=config)

It's not very well documented, but the message definitions are the place to start if you want to do more than that.

Occupies all GPU memory with the following tensorflow code?

1 Answers1