我正在使用 MPI4PY 将 n/p 列分散到两个输入数据进程。但是,我无法按照我的意愿发送专栏。我必须对代码进行哪些更改才能获得最终评论中报告的结果?
矩阵为:
[1, 2, 3, 4]
[5, 6, 7, 8]
[9, 10, 11, 12]
[13, 14, 15, 16]
那么,n=4,p=2。每个进程分别有 2 列。
这是我的代码:
# Imports
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
rows = 4
num_columns = rows/size
data=None
if rank == 0:
data = np.matrix([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]])
recvbuf = np.empty((rows, int(num_columns)), dtype='int')
comm.Scatterv(data, recvbuf, root=0)
print('Rank: ',rank, ', recvbuf received:\n ',recvbuf)
我得到以下输出:
Rank: 0 , recvbuf received:
[[1 2]
[3 4]
[5 6]
[7 8]]
Rank: 1 , recvbuf received:
[[ 9 10]
[11 12]
[13 14]
[15 16]]
我想得到以下输出:
Rank: 0 , recvbuf received:
[[1 2]
[5 6]
[9 10]
[13 14]]
Rank: 1 , recvbuf received:
[[ 3 4]
[7 8]
[11 12]
[15 16]]
我认为这段代码可以满足您的需求。这里的问题是 Scatterv 根本不关心 numpy 数组的形状,它只考虑包含您的值的线性内存块。因此,最简单的方法是预先将数据操作为正确的顺序。请注意,
send_data
是一个一维数组,但这并不重要,因为 Scatterv 并不关心。在另一端,recvbuf 的形状已经定义,Scatterv 只是从收到的 1D 输入中填充它。
# Imports
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
rows = 4
num_cols = rows/size
send_data=None
if rank == 0:
data = np.matrix([[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16]])
# Split into sub-arrays along required axis
arrs = np.split(data, size, axis=1)
# Flatten the sub-arrays
raveled = [np.ravel(arr) for arr in arrs]
# Join them back up into a 1D array
send_data = np.concatenate(raveled)
recvbuf = np.empty((rows, int(num_cols)), dtype='int')
comm.Scatterv(send_data, recvbuf, root=0)
print('Rank: ',rank, ', recvbuf received:\n ',recvbuf)
我使用
send
和 recv
得到的答案略有不同。由于我并不总是得到一个可以均匀分布的矩阵,因此我允许在每个进程中使用不同的切片大小。这对于 Scatterv
来说也是可能的,但我认为使用 send
和 recv
有时更容易处理:
import numpy as np
from mpi4py import MPI
comm = MPI.COMM_WORLD
Np = comm.Get_size()
p = comm.Get_rank()
borders = np.array([0, 1, 2, 3, 5])
if p==0:
arr = np.array([[1, 2, 3, 4, 5],
[6, 7, 8, 9, 0],
[4, 2, 3, 4, 5],
[6, 7, 8, 9, 1],
[2, 2, 3, 4, 5]])
for i in range(1, Np):
ps, pe = borders[i], borders[i+1]
comm.send(arr[ps:pe], dest=i, tag=1)
ps, pe = borders[0], borders[1]
arr = arr[ps:pe]
else:
arr = comm.recv(source=0, tag=1)
print(p, arr)
之前的响应是很好的起点,但下面的
scatter_nd
函数是一个通用的解决方案,因此:
test_scatter_nd.py
import numpy as np
from mpi4py import MPI
def scatter_nd(array, axis=0, comm=MPI.COMM_WORLD, root=0):
"""Scatter n-dimensional array from root to all ranks"""
ans = None
if comm.rank == root:
splits = np.array_split(array, comm.size, axis=axis)
for i in range(comm.size):
if i == root:
ans = splits[i]
else:
comm.send(splits[i], dest=i)
else:
ans = comm.recv(source=root)
return ans
arr = None
if MPI.COMM_WORLD.rank == 0:
arr = np.array([[1, 2, 3, 4, 5],
[6, 7, 8, 9, 10],
[11, 12, 13, 14, 15],
[16, 17, 18, 19, 20],
[21, 22, 23, 24, 25]])
arr = scatter_nd(arr, axis=1)
print(arr, "shape:", arr.shape, "rank:", MPI.COMM_WORLD.rank)
执行
mpiexec -n 2 python test_scatter_nd.py
返回:
[[ 1 2 3]
[ 6 7 8]
[11 12 13]
[16 17 18]
[21 22 23]] shape: (5, 3) rank: 0
[[ 4 5]
[ 9 10]
[14 15]
[19 20]
[24 25]] shape: (5, 2) rank: 1