如何加快张量流的训练速度?

时间:2016-12-12 08:36:20

标签: tensorflow lstm

现在基础tensorflow-char-rnn我开始一个word-rnn项目来预测下一个单词。 但我发现我的火车数据集的速度太慢了。这是我的培训细节:

  • 培训数据量:10亿字
  • 词汇量:0.75百万
  • RNN型号:lstm
  • RNN层:2
  • 细胞大小:200
  • Seq长度:20
  • 批量大小:40(太大的批量大小将导致OOM异常)

机器细节:

  • 亚马逊p2实例
  • 1核K80 GPU
  • 16G视频内存
  • 4核CPU
  • 60G内存

在我的测试中,训练数据1个时代的时间需要17天! 这真的太慢了​​,然后我将seq2seq.rnn_decoder更改为tf.nn.dynamic_rnn,但时间仍然是17天。

我想找到太慢的原因是由我的代码引起的,或者它总是那么慢? 因为我听到一些传闻说Tensorflow rnn比其他DL Framework慢。

这是我的型号代码:

class SeqModel():
def __init__(self, config, infer=False):
    self.args = config
    if infer:
        config.batch_size = 1
        config.seq_length = 1

    if config.model == 'rnn':
        cell_fn = rnn_cell.BasicRNNCell
    elif config.model == 'gru':
        cell_fn = rnn_cell.GRUCell
    elif config.model == 'lstm':
        cell_fn = rnn_cell.BasicLSTMCell
    else:
        raise Exception("model type not supported: {}".format(config.model))

    cell = cell_fn(config.hidden_size)

    self.cell = cell = rnn_cell.MultiRNNCell([cell] * config.num_layers)

    self.input_data = tf.placeholder(tf.int32, [config.batch_size, config.seq_length])
    self.targets = tf.placeholder(tf.int32, [config.batch_size, config.seq_length])
    self.initial_state = cell.zero_state(config.batch_size, tf.float32)

    with tf.variable_scope('rnnlm'):
        softmax_w = tf.get_variable("softmax_w", [config.hidden_size, config.vocab_size])
        softmax_b = tf.get_variable("softmax_b", [config.vocab_size])

        embedding = tf.get_variable("embedding", [config.vocab_size, config.hidden_size])
        inputs = tf.nn.embedding_lookup(embedding, self.input_data)


    outputs, last_state = tf.nn.dynamic_rnn(cell, inputs, initial_state=self.initial_state)

    # [seq_size * batch_size, hidden_size]
    output = tf.reshape(tf.concat(1, outputs), [-1, config.hidden_size])

    self.logits = tf.matmul(output, softmax_w) + softmax_b
    self.probs = tf.nn.softmax(self.logits)

    self.final_state = last_state


    loss = seq2seq.sequence_loss_by_example([self.logits],
                                            [tf.reshape(self.targets, [-1])],
                                            [tf.ones([config.batch_size * config.seq_length])],
                                            config.vocab_size)
    self.cost = tf.reduce_sum(loss) / config.batch_size / config.seq_length

    self.lr = tf.Variable(0.0, trainable=False)
    tvars = tf.trainable_variables()
    grads, _ = tf.clip_by_global_norm(tf.gradients(self.cost, tvars),
                                      config.grad_clip)
    optimizer = tf.train.AdamOptimizer(self.lr)
    self.train_op = optimizer.apply_gradients(zip(grads, tvars))

Here is the GPU load during the training

非常感谢。

3 个答案:

答案 0 :(得分:5)

正如你提到的那样 batch_size 非常重要,它可以带来令人印象深刻的加速,但要检查你的困惑是否具有相关性。

监控GPU活动可以为您提供有关潜在I / O瓶颈的提示。

最重要的是,使用采样softmax 代替常规softmax更快。这需要您使用[config.vocab_size, config.hidden_size]权重矩阵而不是[config.hidden_size, config.vocab_size]。这绝对是我的观点。

希望这有帮助。

pltrdy

答案 1 :(得分:3)

另一种可以加快培训速度的方法,以及缺乏GPU利用率的可能原因是您正在使用占位符。你应该使用队列,如果使用Tensorflow< 1.2,和tf.contrib.data模块否则。

https://www.tensorflow.org/programmers_guide/threading_and_queues

答案 2 :(得分:0)

这里有两行代码加快了我的执行速度。

tf.compat.v1.disable_eager_execution()
tf.config.optimizer.set_jit(True)

请参阅 here 以获取急切执行,以及 here 获取 jit 以判断它是否对您的情况有所帮助。