Is there any way to share variable in Grappler pass?

TensorFlow General Discussion
, ,
1

Hi there,
I want to replica some nodes on different devices in Grappler optimization pass. In this situation, there are some cases when I need to share the Variable node on different devices.

Let’s say I have a node x on GPU :0 , it’s type is an _Arg node. I want to replicate the same node on GPU:1, let’s call it x1. How could I share this variable on graph rewrite pass? If I directly replicate the x node and keep everything but the name the same, it will get an error of something like variable is not initialized.

In addition, tf1 could get a graph that how variable is assigned , it have a graph with Variable/Assign node, but tf2 seems delete this, how could I get how variable is initialized and passed to device?

3

Hello @Stonepia

Thank you for using TensorFlow
In the TensorFlow 2.x, We can create variable by using identity method

with tf.device('/gpu:0'):
    x = tf.Variable(tf.random.normal([10]), name='shared_var')

with tf.device('/gpu:1'):
    x1 = tf.identity(x)

The variable handling is simple because of eager execution, and we can replicate variables across devices easily, without disturbing graph structure.
ps: notes