Advertisement

tensorflow 多GPU训练

阅读量：

在进行多GPU训练时，输入数据的规模会扩展为batch_size乘以GPU数量，从而显著缩短模型的训练时长。

在TensorFlow框架中，若需指定特定GPU进行运算，可通过调用tf.device()函数实现。例如，若希望使用编号为0的显卡进行计算：

复制代码 
    gpu_ind=0
    with tf.device("/gpu:{}".format(gpu_ind))

以下将对多GPU模型训练进行简要说明。相关代码来源于OpenSeq2Seq:https://github.com/NVIDIA/OpenSeq2Seq

针对多GPU模型，其定义文件位于OpenSeq2Seq/model/model_base.py

首先需要设定输入数据，并将其分配至各个GPU的输入端口：

复制代码 
    # placeholders for feeding data
    self.x = tf.placeholder(tf.int32, [self.global_batch_size, None])
    self.x_length = tf.placeholder(tf.int32, [self.global_batch_size])
    self.y = tf.placeholder(tf.int32, [s

全部评论 (0)

还没有任何评论哟~

相关文章推荐