在gpu上运行tensorflow textsum模型

Question

在cpu上，文本和训练代码(https://github.com/tensorflow/models/tree/master/research/textsum)可以完美地运行。

为了在gpu上运行代码以进行更快的训练，我将以下4行中的"cpu“改为"gpu”：https://github.com/tensorflow/models/blob/master/research/textsum/seq2seq_attention.py#L84 https://github.com/tensorflow/models/blob/master/research/textsum/seq2seq_attention_model.py#L149 https://github.com/tensorflow/models/blob/master/research/textsum/seq2seq_attention_model.py#L217 https://github.com/tensorflow/models/blob/master/research/textsum/seq2seq_attention_model.py#L228

模型仍在接受训练，但还是在cpu上，而不是在gpu上。(我用nvidia-smi检查过了)

如果需要在gpu上训练它，请让我知道其他需要更改的地方。

Answer 1

更改这几行代码将帮助您在GPU上训练模型

textsum/seq2seq_attention.py

# change this line
# tf.app.flags.DEFINE_integer('num_gpus', 0, 'Number of gpus used.')
tf.app.flags.DEFINE_integer('num_gpus', 1, 'Number of gpus used.')

textsum/seq2seq_attention_model.py

# if self._num_gpus > 1:
#     self._cur_gpu = (self._cur_gpu + 1) % (self._num_gpus-1)
self._cur_gpu = (self._cur_gpu + 1) % (self._num_gpus-1 if self._num_gpus > 1 else self._num_gpus)

参考问题link