服务器部署多GPU上多个TensorFlow模型,设备/图冲突

根据任务的需要,在服务器后台部署多个GPU模型时,容易出现设备迁移错误的问题。

主要原因还是多个模型的变量在声明时没有新建图或者使用默认图:

一个TensorFlow的运算,被表示为一个数据流的图。

一幅图中包含一些操作(Operation)对象,这些对象是计算节点。前面说过的Tensor对象,则是表示在不同的操作(operation)间的数据节点

一旦任务开始,就已经有一个默认的图已经创建好了。而且可以通过调用:

self.graph = tf.get_default_graph()

这条语句调用系统默认图,在进行简单模型复现训练的时候非常好用,但是在多模型多GPU,尤其是需要将不同模型显式放在不同GPU上时,可能会因为一张默认图中包括不同GPU中的变量而引起的混乱,使用同一BackBone的模型也可能会有变量冲突。

采用的解决方案是:每个模型尽量在声明变量前定义一张图,然后将图设为默认。默认的计算图是和当前的线程绑定的,如果需要在新的线程使用默认的计算图,就要显式的在其他线程通过 with g.as_default(): 进行添加:

self.graph = tf.Graph()
with self.graph.as_default(),tf.device('/gpu:' + device):
	通过上下文来约束图变量和设备,在进行预测时也别忘了指定图:
with self.graph.as_default():
    pred = self.gen.predict(poses_2d)
	为了避免有些不适合GPU上运行的计算,别忘了:
self.sess=tf.Session(graph=self.graph,
		config=tf.ConfigProto(allow_soft_placement=True) )

以下的代码效果相等:

# 1. Using Graph.as_default():
  g = tf.Graph() 
  with g.as_default():
    c = tf.constant(5.0) assert c.graph is g

# 2. Constructing and making default: 
  with tf.Graph().as_default() as g:
    c = tf.constant(5.0) assert c.graph is g

你可能感兴趣的:(服务器部署多GPU上多个TensorFlow模型,设备/图冲突)