tensorflow变量共享——单机多GPU下的参数共享/复用

写作本文的起因，是我作为一个新手，在看单机多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$、$hiddenstate$、$cellstate$）都是每个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行）是为了在不同的设备之间复用某些网络（例如对于$fact$和$question$进行embedding的网络和进行encode的网络）的参数。