2017-09-06 7 views
0

私はword2vecを使ってコードを実験していましたhttps://github.com/chiphuyen/stanford-tensorflow-tutorials/blob/master/examples/04_word2vec_no_frills.py次のテンソルフローコードですべてのGPUメモリを占有しますか?

しかし、私はすべてのGPUメモリを簡単に使い切っています。

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() 

memory used up

+0

TFの標準的な動作だと思います。 [ここ](https://stackoverflow.com/questions/34199233/how-to-prevent-tensorflow-from-allocating-the-totality-of-a-gpu-memory)を参照してください。たとえばgoogleの場合、「tensorflowはすべてのGPUメモリを割り当てます」という説明があります。 –

+0

@RobertCrovellaは驚いた! –

答えて

4

これがデフォルトtensorflowの動作です。 GPUのメモリ割り当てを必要なものだけに制限する場合は、セッション設定でこれを指定します。

代わりに使用するGPUメモリの最大割合を指定することができます。

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

それは非常によく文書化されていますが、それ以上のことをしたい場合はmessage definitionsを開始する場所はないのです。

関連する問題