模型尺寸太大,我的注意力模型实现如何?

时间:2019-04-09 09:39:31

标签: tensorflow machine-translation attention-model

我正在实施Minh-Thang Luong的注意力模型来构建英语到中文的翻译器。我训练的模型大小异常大(980 MB)。Minh-Thang Luong's original paper
model structure

这是模型参数

state size:120  
source language vocabulary size:400000  
source language word embedding size:400000*50  
target language vocabulary size:20000  
target language word embedding size:20000*300

这是我在tensorflow中的模型实现。

import tensorflow as tf

src_vocab_size=400000
src_w2v_dim=50
tgt_vocab_size=20000
tgt_w2v_dim=300
state_size=120

with tf.variable_scope('net_encode'):
    ph_src_embedding = tf.placeholder(dtype=tf.float32,shape=[src_vocab_size,src_w2v_dim],name='src_vocab_embedding_placeholder')
    #src_word_emb = tf.Variable(initial_value=ph_src_embedding,dtype=tf.float32,trainable=False, name='src_vocab_embedding_variable')

    encoder_X_ix = tf.placeholder(shape=(None, None), dtype=tf.int32)
    encoder_X_len = tf.placeholder(shape=(None), dtype=tf.int32)
    encoder_timestep = tf.shape(encoder_X_ix)[1]
    encoder_X = tf.nn.embedding_lookup(ph_src_embedding, encoder_X_ix)
    batchsize = tf.shape(encoder_X_ix)[0]

    encoder_Y_ix = tf.placeholder(shape=[None, None],dtype=tf.int32)
    encoder_Y_onehot = tf.one_hot(encoder_Y_ix, src_vocab_size)

    enc_cell = tf.nn.rnn_cell.LSTMCell(state_size)
    enc_initstate = enc_cell.zero_state(batchsize,dtype=tf.float32)
    enc_outputs, enc_final_states = tf.nn.dynamic_rnn(enc_cell,encoder_X,encoder_X_len,enc_initstate)
    enc_pred = tf.layers.dense(enc_outputs, units=src_vocab_size)
    encoder_loss = tf.losses.softmax_cross_entropy(encoder_Y_onehot,enc_pred)
    encoder_trainop = tf.train.AdamOptimizer(0.001).minimize(encoder_loss)

with tf.variable_scope('net_decode'):
    ph_tgt_embedding = tf.placeholder(dtype=tf.float32, shape=[tgt_vocab_size, tgt_w2v_dim],
                                      name='tgt_vocab_embedding_placeholder')
    #tgt_word_emb = tf.Variable(initial_value=ph_tgt_embedding, dtype=tf.float32, trainable=False, name='tgt_vocab_embedding_variable')
    decoder_X_ix = tf.placeholder(shape=(None, None), dtype=tf.int32)
    decoder_timestep = tf.shape(decoder_X_ix)[1]
    decoder_X_len = tf.placeholder(shape=(None), dtype=tf.int32)
    decoder_X = tf.nn.embedding_lookup(ph_tgt_embedding, decoder_X_ix)

    decoder_Y_ix = tf.placeholder(shape=[None, None],dtype=tf.int32)
    decoder_Y_onehot = tf.one_hot(decoder_Y_ix, tgt_vocab_size)

    dec_cell = tf.nn.rnn_cell.LSTMCell(state_size)
    dec_outputs, dec_final_state = tf.nn.dynamic_rnn(dec_cell,decoder_X,decoder_X_len,enc_final_states)

    tile_enc = tf.tile(tf.expand_dims(enc_outputs,1),[1,decoder_timestep,1,1]) # [batchsize,decoder_len,encoder_len,state_size]
    tile_dec = tf.tile(tf.expand_dims(dec_outputs, 2), [1, 1, encoder_timestep, 1]) # [batchsize,decoder_len,encoder_len,state_size]
    enc_dec_cat = tf.concat([tile_enc,tile_dec],-1) # [batchsize,decoder_len,encoder_len,state_size*2]
    weights = tf.nn.softmax(tf.layers.dense(enc_dec_cat,units=1),axis=-2) # [batchsize,decoder_len,encoder_len,1]
    weighted_enc = tf.tile(weights, [1, 1, 1, state_size])*tf.tile(tf.expand_dims(enc_outputs,1),[1,decoder_timestep,1,1]) # [batchsize,decoder_len,encoder_len,state_size]
    attention = tf.reduce_sum(weighted_enc,axis=2,keepdims=False) # [batchsize,decoder_len,state_size]
    dec_attention_cat = tf.concat([dec_outputs,attention],axis=-1) # [batchsize,decoder_len,state_size*2]
    dec_pred = tf.layers.dense(dec_attention_cat,units=tgt_vocab_size) # [batchsize,decoder_len,tgt_vocab_size]
    pred_ix = tf.argmax(dec_pred,axis=-1) # [batchsize,decoder_len]
    decoder_loss = tf.losses.softmax_cross_entropy(decoder_Y_onehot,dec_pred)
    total_loss = encoder_loss + decoder_loss
    decoder_trainop = tf.train.AdamOptimizer(0.001).minimize(total_loss)

_l0 = tf.summary.scalar('decoder_loss',decoder_loss)
_l1 = tf.summary.scalar('encoder_loss',encoder_loss)
log_all = tf.summary.merge_all()
writer = tf.summary.FileWriter(log_path,graph=tf.get_default_graph())

这是到目前为止我能想到的模型参数大小的精简版

encoder cell
=(50*120+120*120+120)*4
=(src_lang_embedding_size*statesize+statesize*statesize+statesize)*(forget gate,remember gate,new state,output gate)
=(kernelsize_for_input+kernelsize_for_previous_state+bias)*(forget gate,remember gate,new state,output gate)  
=82080 floats

encoder dense layer  
=120*400000
=statesize*src_lang_vocabulary_size
=48000000 floats

decoder cell
=(300*120+120*120+120)*4
=(target_lang_embedding_size*statesize+statesize*statesize+statesize)*(forget gate,remember gate,new state,output gate)
=(kernelsize_for_input+kernelsize_for_previous_state+bias)*(forget gate,remember gate,new state,output gate)
=202080 floats

dense layer that compute attention weights
=(120+120)*1
=(encoder_output_size+decoder_output_size)*(1 unit)
=240 floats

decoder dense layer
=(120+120)*20000
=(attention_vector_size+decoder_outputsize)*target_lang_vocabulary_size
=4800000 floats

它们全部获得212 MB,但实际模型大小为980 MB。那么哪里出了错?

1 个答案:

答案 0 :(得分:1)

您仅计算可训练参数的数量,这些并不是您需要容纳在GPU内存中的唯一数字。

您正在使用Adam优化器,因此,您需要存储所有参数的梯度和所有参数的动量。这意味着您需要将每个参数存储3次,这将为您提供636 MB。

然后,您需要容纳网络的所有中间状态以进行正向和反向传递。

比方说,批次大小为 b ,来源为目标长度为50,那么您就有了(至少,我可能忘记了一些东西):

  • b × l ×50个源嵌入,
  • b × l ×300个目标嵌入,
  • b × l ×5×120编码器状态,
  • b × l ×400000编码器对数,
  • b × l ×5×300解码器状态,
  • b × l ×120个中间注意状态,
  • b × l ×20000输出logits。

总共需要存储421970× b × l 浮点数。

顺便说一句。来源词汇表400k数量巨大,我相信其中大多数都不足够频繁地学习关于它们的任何有意义的知识。您应该使用预处理程序(即SentencePiece),以将词汇量减少到合理的大小。