且构网

分享程序员开发的那些事...
且构网 - 分享程序员编程开发的那些事

tensorflow中模型并行的实现

更新时间:2023-12-02 19:24:28

这是一个例子.该模型有一部分在GPU0上,一部分在GPU1上,一部分在CPU上,所以这是3路模型并行.

Here's an example. The model has some parts on GPU0, some parts on GPU1 and some parts on CPU, so this is 3 way model parallelism.

with tf.device("/gpu:0"):
    a = tf.Variable(tf.ones(()))
    a = tf.square(a)
with tf.device("/gpu:1"):
    b = tf.Variable(tf.ones(()))
    b = tf.square(b)
with tf.device("/cpu:0"):
    loss = a+b
opt = tf.train.GradientDescentOptimizer(learning_rate=0.1)
train_op = opt.minimize(loss)

sess = tf.Session()
sess.run(tf.global_variables_initializer())
for i in range(10):
    loss0, _ = sess.run([loss, train_op])
    print("loss", loss0)