Python torch：最小填充张量，使得num元素可被x整除

假设我有一个任意ndim的张量

t

我想（用零）填充它，以便 a） 我引入了尽可能少的元素 b） 填充后，

（t.numel（）%x）==0

有比这更好的算法吗 找到最大尺寸并将其增加1，直到满足条件（b）

可能是工作代码：

def pad_minimally(t, x):
    largest_dim = np.argmax(t.shape)
    buffer_shape = list(t.shape)
    new_t = t.clone()
    print(t.shape)
    for n_to_add in range(x):
        if new_t.numel() % x == 0:
            break
        buffer_shape[largest_dim] = n_to_add
        new_buffer = torch.zeros(*buffer_shape)
        new_t = torch.cat([t, new_buffer], axis=largest_dim)
    assert new_t.numel() % x == 0
    return new_t
assert pad_minimally(torch.rand(3,1), 7).shape == (7,1)
assert pad_minimally(torch.rand(3,2), 7).shape == (7,2)
assert pad_minimally(torch.rand(3,2, 6), 7).shape == (3,2,7)

---

首先，简单地向最大维度添加一个，直到

numel

可被

x

整除，这在所有情况下都不起作用。例如，如果

t

的形状是

（3,2）

和

x=9

，那么我们希望将

t

填充为

（3,3）

，而不是

（9,2）

更令人担忧的是，不能保证只需要填充一个维度。例如，如果

t

具有形状

（13,17,25）

和

x=8

，则最佳填充

t

将是

（14,18,26）

或

（13,18,28）

把它提炼成数学，问题就变成了

给定正整数

s[1]，…，s[D]

查找正整数

q[1]，…，q[D]

使

prod（q[i]，i=1到D）

最小化

prod（q[i]，i=1到D）

对于所有

i=1到D

的约束条件是

prod（q[i]，i=1到D）

我没能开发出一个有效的解决方案（参见更新以获得更有效的解决方案），尽管我对非线性整数规划不是特别精通。也许这个问题存在一个有效的解决办法。如果是这样的话，我想它会涉及到

x

和

q

和/或更好的记忆。也就是说，如果

x

和

D

（即

len（t.shape）

）足够小（否则该算法可能运行很长时间），则可以使用穷举搜索解决该问题

我提出的蛮力搜索算法对大于或等于

t.numel（）

的

x

的每一个倍数进行迭代，并使用深度优先搜索查看该倍数是否存在填充。一旦找到有效的填充，算法就会完成。此算法的python代码为：

将numpy导入为np
def搜索（形状、目标、内存）：
numel=np.prod（形状）
如果numel==target\u numel：
返回真值
elif numel

一旦有了最小的形状，

pad_minimal

函数就可以非常简洁地实现为

def pad_最小值（t，x）：
新形状=最小形状（t形，x）
new_t=t.new_零（new_形状）
新的_t[[t形中s的切片（0，s）]]=t
返回新的

我不确定这是否能满足你的需要。希望其他人能提供一个更高效的版本

---

最小_形的一些测试用例

断言最小_形（[2,2]，9）=[3,3]
断言最小_形（[2,8]，6）=[2,9]
在[14,18,26]，[13,18,28]中断言最小_形（[13,17,25]，8）]
断言最小_形（[5,13,19]，6）=[5,14,21]

---

更新 我是CS.SE。根据我在那里得到的答案以及随后对该问题的更新，以下是

minimal\u shape

的一个更有效的实现

from functools import reduce
from operator import mul
from copy import deepcopy

def prod(x):
    return reduce(mul, x, 1)

def argsort(x, reverse=False):
    return sorted(range(len(x)), key=lambda idx: x[idx], reverse=reverse)

def divisors(v):
    """ does not include 1 """
    d = {v} if v > 1 else set()
    for n in range(2, int(v**0.5) + 1):
        if v % n == 0:
            d.add(n)
            d.add(v // n)
    return d

def update_memory(b, c_rem, memory):
    tuple_m = tuple(b + [c_rem])
    if tuple_m in memory:
        return False
    memory.add(tuple_m)
    return True

def dfs(a, b, c, c_rem, memory, p_best=float('inf'), b_best=None):
    ab = [ai + bi for ai, bi in zip(a, b)]
    p = prod(ab)
    if p >= p_best:
        return p_best, b_best
    elif p % c == 0:
        return p, deepcopy(b)

    dc = divisors(c_rem)
    for i in argsort(ab):
        for d in dc:
            db = (d - ab[i]) % d
            b[i] += db
            if update_memory(b, c_rem // d, memory):
                p_best, b_best = dfs(a, b, c, c_rem // d, memory, p_best, b_best)
            b[i] -= db

    return p_best, b_best

def minimal_shape(shape, target_multiple):
    a = list(shape)
    b = [0 for _ in range(len(a))]
    c = target_multiple
    _, b = dfs(a, b, c, c, set())
    return [ai + bi for a, b in zip(a, b)]

---

您提出的贪婪方法也并不总是导致最小的填充。例如，如果您的原始形状是

（3，2）

，并且您希望numel可以被9整除，那么最小的填充将使您的张量成为

（3，3）

，而您的算法将产生

（9，2）

。