multi task训练torch_使用tensorflow1.x实现单机多卡训练

前言

tensorflow1.x关于如何使用单机多卡训练仅有一个官方repo的脚本参考

对于刚接触的小伙伴来说还是有些难懂，这里介绍下tensorflow单机多卡的原理及实现(尽量用最简短的代码)。仅作参考，如有错误及不足之处，望指正！

使用tensorflow2的朋友可以参考官方文档

正文

实验环境

tensorflow-gpu=1.13
python=3.6

单机多卡示意图

首先看一下单机单卡的整个流程(看不清就尝试放大，我已经尽力了):

从输入到损失不必多说，最后compute_gradients,再apply_grandients.一般1.x版本的书上或者教程上都是用的二者合并的接口optimizer.minimize(total_loss)。但是单机多卡的时候得拆开用，下面是单机多卡、模型并行的示意图(假设是3块gpu)：

总结来说就是：每块gpu上都有个同样的模型net，他们权重变量共享，在图中表示为(var0,var1,...,varn)，tf在每块gpu上各自求gradients，然后对每块gpu上的每个var的gradients求和求平均，最后apply_gradients.

核心代码示例一

在github上大部分单机多卡的实现也是如此，下面代码可以作为一个模板参考。

import tensorflow as tf
import os
os.environ['CUDA_VISIBLE_DEVICES']='0,6,7' #