怎么把一个大数组的数据尽量平均地分成几组？

静待花开

多GPU处理同一个大型数组时，我想尽量地把这个数组的数据平均分成几个部分，每个GPU负责处理一部分，但是因为这个数组的大小并不确定，并不一定是GPU数量的整数倍，所以直接相除并不总得到整数。不知道大家有没有遇到类似这种问题，可否提供一下思路，或者参考代码就可以。

li913

没要求非得平均，近似平均就行。

fcode

让其中一个GPU受累多算点。

静待花开

fcode 发表于 2024-6-19 10:13
让其中一个GPU受累多算点。

哈哈，想起来了chatgpt，下面是它的回答。

def divide_integer(n, k):
    # 计算基础值和余数
    q = n // k
    r = n % k
   
    # 初始化结果数组，每个元素为 q
    parts = [q] * k
   
    # 分配余数给前 r 个元素
    for i in range(r):
        parts += 1
   
    return parts

# 示例
n = 10
k = 3
result = divide_integer(n, k)
print(result)  # 输出: [4, 3, 3]

静待花开

li913 发表于 2024-6-19 09:44
没要求非得平均，近似平均就行。

是的，尽量做得平均，关键是每个GPU上要分配的数量是正整数。哈哈