Tensorflow 多 GPU 训练

介绍

TensorFlow中的并行主要分为模型并行和数据并行。 模型并行需要根据不同模型设计不同的并行方式， 其主要原理是将模型中不同计算节点放在不同硬件资源上运算。 比较通用的且能简便地实现大规模并行的方式是数据并行， 其思路我们在第1章讲解过， 是同时使用多个硬件资源来计算不同batch的数据的梯度， 然后汇总梯度进行全局的参数更新。

数据并行又分为同步和异步，同步训练是指等所有GPU得到梯度后统一更新梯度，异步训练是每当有GPU得到梯度就更新网络权重。

这里使用数据并行的同步训练方法，大致如下：

1. 先定义 batch 数，学习率衰减，优化方法等
2. 分别为机器上所有的GPU定义模型，得到损失及梯度，使用reuse_variables，它们共享网络的参数值
3. 保存所有GPU的梯度，得到平均值
4. 根据定义的优化方法及平均梯度，优化模型

示例代码

import os.path
import re
import time
import numpy as np
import tensorflow as tf
import cifar10


batch_size = 128
max_steps = 1000000
num_gpus = 4


def tower_loss(scope):
    '''计算单个设备的所有损失
    
    参数：
        scope：设备的唯一编号
        
    返回值：
        批数据损失的tensor
    '''
    
    image, labels = cifar10.distorted_inputs()
    # 构建模型，获取logit
    logits = cifar10.inference(image)
    
    _ = cifar10.loss(logits, labels)
    # 搜集当前设备的损失
    losses = tf.get_collection('losses', scope)
    # 计算总损失
    total_loss = tf.add_n(losses, name='total_loss')
    
    return total_loss

def average_gradients(tower_grads):
    '''将搜集到的所有设备的梯度 tower_grads 取平均值后作为网络更新梯度
    
    参数： 
        tower_grads: shape 为 [[(grad0_gpu0,var0_gpu0),(grad1_gpu0,var1_gpu0),...],[...]]
        第一个维度为GPU各自的梯度
        第二个维度为GPU所有的变量梯度
        第三个维度是一个tuple，保存参数梯度grad0_gpu0，及参数值var0_gpu0
        
    返回值：
        所有GPU梯度的平均值及变量值，形状为输入参数的前2个维度
        
    操作：
        1、将梯度转置，方便获取所有GPU的同一个变量的同一个梯度
        2、将梯度拓展一个维度，加入list中
        3、合并列表取平均值
        4、获取变量var，与平均梯度一起返回
        
    '''
    average_grads = []
    for grad_and_vars in zip(*tower_grads):
        grads = []
        for g, _ in grad_and_vars:
            expanded_g = tf.expand_dims(g, 0)
            grads.append(expanded_g)
        
        grad = tf.concat(grads, 0)
        grad = tf.reduce_mean(grad, 0)
        
        v = grad_and_vars[0][1]
        grad_and_var = (grad, v)
        
        average_grads.append(grad_and_var)
        
        return average_grads


def train():
    '''训练函数
    
    步骤：
    1、先使用CPU进行一些简单的计算，比如global_step，学习率的衰减
    2、对每一个GPU，使用tower_loss计算自己的损失
    3、调用tf.get_variable_scope().reuse_variables()，使得所有GPU使用同样的参数
    4、计算平均损失，并更新
    
    疑问：
    1、网络是怎样在GPU之间复制的        tower_loss 函数已经创建网络并定义损失
    2、每个GPU怎样获取到自己的数据       tower_loss 函数里包含获取批次数据代码
    3、每个GPU怎样更新网络的权重        多个GPU共用一个Sess，因此使用一样的apply_gradients即可
    4、每个GPU怎样保持一样的权重        reuse_variables 函数使得GPU之间复用权重
    5、start_queue_runners 是做什么用的 准备好大量的数据增强后的样本
    '''
    
    # 全局训练步数
    with tf.Graph().as_default(), tf.device('/cpu:0'):
        global_step = tf.get_variable('global_step', [], 
                                      initializer=tf.constant_initializer(0),
                                      trainable=False)
        
        num_batches_per_epoch = cifar10.NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN / batch_size
        decay_step = int(num_batches_per_epoch * cifar10.NUM_EPOCHS_PER_DECAY)
        
        # 衰减学习率计算
        lr = tf.train.exponential_decay(cifar10.INITIAL_LEARNING_RATE,
                                        global_step,
                                        decay_step,
                                        cifar10.LEARNING_RATE_DECAY_FACTOR,
                                        staircase=True)
        
        opt = tf.train.GradientDescentOptimizer(lr)
        
        # 对每个GPU 使用tower_loss 构建模型，计算损失
        tower_grads = []
        for i in range(num_gpus):
            with tf.device('/gpu:%d' % i):
                with tf.name_scope('%s_%d' % (cifar10.TOWER_NAME, i)) as scope:
                    loss = tower_loss(scope)
                    # GPU共享参数
                    tf.get_variable_scope().reuse_variables() 
                    grads = opt.compute_gradients(loss)
                    tower_grads.append(grads)
            
        # 获取平均梯度并更新
        grads = average_gradients(tower_grads)
        apply_gradient_op = opt.apply_gradients(grads, global_step=global_step)
        
        # 初始化变量，创建Session
        saver = tf.train.Saver(tf.all_variables())
        init = tf.global_variables_initializer()
        sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
        sess.run(init)
        tf.train.start_queue_runners(sess=sess) # ?
        
        # 开始训练并保存数据
        for step in range(max_steps):
            start_time = time.time()
            
            _, loss_value = sess.run([apply_gradient_op, loss])
            duration = time.time() - start_time
            
            if step % 10 == 0:
                num_examples_per_step = batch_size * num_gpus
                examples_per_sec = num_examples_per_step / duration
                sec_per_batch = duration / num_gpus
                
                format_str = ('step %d, loss = %.2f (%.1f examples/sec; %.3f sec/batch)')
                print(format_str % (step, loss_value, examples_per_sec, sec_per_batch))
                
            if step % 1000 == 0 or (step + 1) == max_steps:
                saver.save(sess, './cifar10_train/model.ckpt', global_step=step)
                
cifar10.maybe_download_and_extract()
train()

步骤：
1、先使用CPU进行一些简单的计算，比如global_step，学习率的衰减
2、对每一个GPU，使用tower_loss计算自己的损失
3、调用tf.get_variable_scope().reuse_variables()，使得所有GPU使用同样的参数
4、计算平均损失，并更新

疑问：
1、网络是怎样在GPU之间复制的        tower_loss 函数已经创建网络并定义损失
2、每个GPU怎样获取到自己的数据       tower_loss 函数里包含获取批次数据代码
3、每个GPU怎样更新网络的权重        多个GPU共用一个Sess，因此使用一样的apply_gradients即可
4、每个GPU怎样保持一样的权重        reuse_variables 函数使得GPU之间复用权重
5、start_queue_runners 是做什么用的 准备好大量的数据增强后的样本

其他比如 cifar10.py 的代码参考这里https://github.com/yhlleo/tensorflow.cifar10

参考

《Tensorflow实战》 9.2 多GPU并行