2016-07-14 77 views
0

我試圖在tensorflow上創建1,000,000個字的嵌入。每個單詞將有一個256 float32矢量表示該單詞。問題在於我一直在耗盡內存。由於我的GTX 1080擁有8GB的內存,因此這對我來說並不重要。嵌入應該只佔用1e6 * 256 * 4 = 1 Gb的內存。我還在輸出上有另一個相同大小的矩陣。除此之外,還有其他一些張量應該比較小。因此,我只能看到存儲該模型所需的大約2 - 3 GB的內存,當我撥打sess.run(tf.initialize_all_variables())時,該內存會失效。我的記憶在哪裏?你對我如何解決這個問題有任何建議嗎?Tensorflow嵌入空間不足

import tensorflow as tf 
import nltk 
import numpy as np 
import os 
import multiprocessing 
import itertools 
import pickle 
from unidecode import unidecode 

BATCH_SIZE = 32 
TIME_STEPS = 64 
WORD_VEC_SIZE = 256 

words, training_data = pickle.load(open('vocab.pickle', 'rb')) 
word2index = {w:i for i, w in enumerate(words)} 
index2word = {i:w for i, w in enumerate(words)} 

input_tensor = tf.placeholder(tf.int32, (BATCH_SIZE, TIME_STEPS + 1), 'input_tensor') 
embedding = tf.Variable(tf.random_uniform((len(words), WORD_VEC_SIZE), -1, 1), name = 'embedding') 

rnn = tf.nn.rnn_cell.BasicRNNCell(WORD_VEC_SIZE) 
state = tf.zeros((BATCH_SIZE, rnn.state_size)) 
input_vectors = tf.nn.embedding_lookup([embedding], input_tensor[:, :TIME_STEPS]) 
cost = 0 

with tf.variable_scope('rnn') as scope: 
    W_out = tf.get_variable('W_out', (WORD_VEC_SIZE, len(words)), initializer = tf.truncated_normal_initializer(0.0, 1/np.sqrt(WORD_VEC_SIZE))) 
    b_out = tf.get_variable('b_out', (len(words),), initializer = tf.truncated_normal_initializer(0.0, 0.01)) 
    for t in range(TIME_STEPS): 
     y, state = rnn(tf.reshape(input_vectors[:, t, :], (-1, WORD_VEC_SIZE)), state) 
     cost += tf.reduce_mean(tf.nn.sampled_softmax_loss(W_out, b_out, y, tf.reshape(input_tensor[:, t + 1], (-1, 1)), 1000, len(words))) 
     scope.reuse_variables() 

train_step = tf.train.AdamOptimizer(1e-4).minimize(cost) 

sess = tf.Session() 
sess.run(tf.initialize_all_variables()) 
saver = tf.train.Saver() 

回答

2

我沒有考慮的是AdamOptimizer。我忘記了這需要爲我的模型中的每個權重存儲各種參數。當我更改爲GraidentDecent優化器時,它現在適合我的GPU。