根据任务的需要,在服务器后台部署多个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