簡體   English   中英

Tensorflow嵌入空間不足

[英]Tensorflow embedding running out of space

我正在嘗試在tensorflow上為1,000,000個單詞創建一個嵌入。 每個單詞都有一個代表該單詞的256 float32向量。 問題是我一直用不完內存。 這對我沒有意義,因為我的GTX 1080上有8GB的內存。嵌入應該只占用1e6 * 256 * 4 = 1 Gb的內存。 我在輸出上還有另一個大小相同的矩陣。 除此之外,還有一些其他張量相比而言應該較小。 因此,我只看到需要大約2-3 GB的內存來存儲模型,並且在調用sess.run(tf.initialize_all_variables())時失敗。 我所有的記憶都去了哪里,您對我如何解決這個問題有任何建議嗎?

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

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

暫無
暫無

聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.

 
粵ICP備18138465號  © 2020-2024 STACKOOM.COM