更新时间: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)