我想在 2 个 GPU 上并行化以下简单表达式:在对结果求和之前在 GPU 0 和GPU 1 上进行C = A^n + B^n
计算。A^n
B^n
在 TensorFlow 中,我会喜欢:
with tf.device('/gpu:0'):
An = matpow(A, n)
with tf.device('/gpu:1'):
Bn = matpow(B, n)
with tf.Session() as sess:
C = sess.run(An + Bn)
但是,由于 PyTorch 是动态的,所以我在做同样的事情时遇到了麻烦。我尝试了以下方法,但只需要更多时间。
with torch.cuda.device(0):
A = A.cuda()
with torch.cuda.device(1):
B = B.cuda()
C = matpow(A, n) + matpow(B, n).cuda(0)
我知道有一个模块可以在批处理维度上并行化模型,torch.nn.DataParallel
但在这里我尝试做一些更基本的事情。