写作本文的起因,是我作为一个新手,在看单机多GPU的tensorflow代码时,看到了一段很费解的代码,完整代码戳这里。因为不懂VariableScope和NameScope的作用和区别,看着这段好多个with的代码觉得非常乱。所以这里记录下自己的分析过程(笔记来的,散了吧):
......
from tensorflow.contrib import layers
from tensorflow.contrib.rnn import LSTMCell
......
with tf.variable_scope(tf.get_variable_scope()):
for device_nr, device in enumerate(self.devices):
with tf.device('/cpu:0'):
......
facts_emb = layers.embed_sequence(facts_ph, self.vocab.size(), self.emb_size, scope='word-embeddings')
questions_emb = layers.embed_sequence(question_ph, self.vocab.size(), self.emb_size, scope='word-embeddings', reuse=True)
with tf.device(device), tf.name_scope("device-%s" % device_nr):
def mlp(x, scope, n_hidden):
with tf.variable_scope(scope):
for i in range(3):
x = layers.fully_connected(x, n_hidden, weights_regularizer=regularizer)
return layers.fully_connected(x, n_hidden, weights_regularizer=regularizer, activation_fn=None)
_, (_, f_encoding) = tf.nn.dynamic_rnn(tf.nn.rnn_cell.LSTMCell(32), facts_emb, dtype=tf.float32, sequence_length=f_seq_length_ph, scope='fact-encoder')
......
_, (_, q_encoding) = tf.nn.dynamic_rnn(tf.nn.rnn_cell.LSTMCell(32), questions_emb, dtype=tf.float32, sequence_length=q_seq_length_ ph, scope='question-encoder')
......
with tf.variable_scope('steps'):
lstm_cell = LSTMCell(self.n_hidden)
state = lstm_cell.zero_state(tf.shape(x)[0], tf.float32)
for step in range(self.n_steps):
x = message_passing(x, edge_indices_ph, edge_features, lambda x: mlp(x, 'message-fn', self.n_hidden), edge_keep_prob)
x = mlp(tf.concat([x, x0], axis=1), 'post-fn', self.n_hidden)
x, state = lstm_cell(x, state)
with tf.variable_scope('graph-sum'):
......
tf.get_variable_scope().reuse_variables()
......
tf.get_variable_scope().reuse_variables()
上面的代码一共有3处是与reuse模式相关的代码,分别是第6行、第29行和第31行。第6行在调用tensorflow.contrib.layers.embedsequence()
时传入了reuse=True
,而29和31行各有一句tf.get_variable_scope().reuse_variables() ,现在分别解释他们的作用。
首先,第5行已经调用了一次layers.embedsequence()
用于求embedding。而我们肯定希望在同一个问题中,对所有的文本数据都使用同一套embedding的参数,所以在第6行再次调用layers.embedsequence()
时,我们与第5行一样,传入scope = word-embeddings
,并且传入reuse = True
,这就告诉tensorflow,我们之前在'word-embeddings'
域下新建的用于训练的参数都要复用/共享。通过查看官方文档对于tensorflow.contrib.layers.embedsequence()
的参数的解释,可以看到对于scope参数的解释是
Optional string specifying the variable scope for the op, required if reuse=True.
进一步查看还可以发现layers里面的很多Operation,例如tf.contrib.layers.fully_connected
,都有scope和reuse参数,而且都是类似的解释。看到这里大家就应该明白scope和reuse这两个参数的含义以及怎么用了吧。
接下来继续分析第29行和第31行的设置reuse模式的作用。
首先我们要知道作者是想对哪个scope设置reuse模式。通过在29行和31行的前一行加上一句print(tf.get_variable_scope().name)
然后运行,即可知道第29行目的是将step
这个变量域以及其子变量域设置为reuse模式,而31行则是将默认变量域(其name为空白字符串)以及其子域设置为reuse模式。
我们先来分析第29行。
在第29行之后,steps
域以及其子域下新建的用于训练的参数都要复用,那么这些参数有哪些呢?第20行是steps
域开始的地方,可以看到在第24、25、26、27行分别在steps
域下设置了4个子域,分别是steps/message_fn
,steps/post_fn
,steps/lstm_cell
,steps/graph_sum
。
其中前两个子域内调用了自定义的mlp
方法,从mlp
的定义(第9-13行)可以发现,其实就是间接调用了一个4层的全连接层,所以第一次运行第24、25行时,这两个由4层全连接网络组成的Multilayer Perceptron (MLP) 的参数是新建的,在之后(for循环的step变量大于0时)这两个MLP的参数都是要复用的。LSTM层的参数复用见后面的题外话,steps/graph_sum
域的分析这里略去,因为与前两个差不多。
+++++++++++++++++题外话:为什么说steps/lstm_cell
是一个VariableScope?begin +++++++++++++++++
通过查看官方文档发现,在第26行定义LSTMCell时,没有传入name
参数,这个参数用于指定Layer的名称,而如果希望相同name
的Layer共享参数(此处是LSTM层),必须同时指定reuse=True
。
Layers with the same name will share weights, but to avoid mistakes we require reuse=True in such cases.
因此个人认为Layer的name是用来指定其内部的Variable的VariableScope的(此处没传入,使用默认的"lstm_cell"),进而,在第29行设置reuse模式后(相当于文档要求的“require reuse=True”),这个LSTM层的参数就可以被复用了。不过我看了半天tensorflow的源码没找到将name
设置为namescope的代码,所以我这样的理解不一定是最准确的,有可能Layer的name并非用来指定其中变量的VariableScope。至少,这个name是用来作为Layer内部Variable的name的前缀的一部分的。总之,记得LSTM层的复用是这样设置即可。
++++++++++++++++++题外话:为什么说steps/lstm_cell
是一个VariableScope?end ++++++++++++++++++
为了验证一下,我加上了一些打印代码,看一下在4个GPU的服务器上跑,会打印出什么。用于打印的主要代码如下,加在第31行前,由于第2行的for循环,一共会打印4次网络的参数:
for v in tf.trainable_variables(): # 打印所有(可训练的)参数,再打印其所在设备。
print( e, e.device )
打印结果的一部分如下,可以清晰地看到steps内部的4个子变量域的所有参数:
第31行的作用与第29行的位置和作用都是类似的,都是放在for的循环体最后,受它影响而被复用的网络参数是下面这些:用于得到embedding的矩阵参数,用于对fact进行encode的lstm(包括kernel和bias),用于对question进行encode的lstm,用于对encode之后的向量进行处理的MLP(4层FC的,得到的结果输入RRN的node)。
这两块输出拼在一起就是整个网络的所有可训练的参数的信息了,不过这只是在GPU0(第2行for循环的第一次运行)上的输出结果。而在其他GPU上的输出结果与第一次输出的结果完全相同,因此可以看出,网络内部的所有trainable的参数都在CPU或第一个GPU上保存着,而后面3个GPU是复用这些参数进行计算。当然,每个GPU上网络的输入和输出(例如LSTM的 x x x、 h i d d e n s t a t e hidden\space state hidden state、 c e l l s t a t e cell\space state cell state)都是每个GPU各一份的,不进行共享(对这些Tensor打印信息即可验证)。
小结
通过在for循环的循环体的末尾部分使用tf.get_variable_scope().reuse_variables()
,可以使得在for循环内调用的网络参数被复用。即第一次运行循环体时tensorflow会新建参数,后面循环时tensorflow只需复用第一次运行循环体时新建的参数。
在本小节的代码中,作者用了2次。
第一次(第29行)是为了在不同时间步之间复用recurrent-relational-network(RRN)内的几个部分参数,包括实体(entity,此任务中定义为fact)之间的消息传递网络message-fn
、每个实体之后接的MLPpost-fn
、LSTM运算lstm_cell
、Loss等的计算graph-sum
。
第二次(第31行)是为了在不同的设备之间复用某些网络(例如对于 f a c t fact fact和 q u e s t i o n question question进行embedding的网络和进行encode的网络)的参数。