Question

我想在优化过程中实施更新规则。我想使用pytorch函数对其向量化，以提高时间复杂度。可能在GPU上运行。

更新规则就像

u[i][j] = multiplier*NORM(Q.column(i)-Q.column(j))

我尝试为其执行编写一个简单的嵌套循环，但是看起来异常缓慢。

def u_vector(Q,parameters):
    u=torch.zeros(Q.shape[0],Q.shape[0],requires_grad=False)
    for i in range(Q.shape[0]):
        for j in range(i+1,Q.shape[1]):
            u[i][j]=0.5*(torch.norm(Q[:,i]-Q[:,j]))*parameters['gamma']
            u[j][i]=u[i][j]
    return u

有人可以帮忙吗？

Answer 1

这是使用索引一次计算所有差异和规范的一种可能性。您需要进行测试，看看它在您的硬件上是否更快。

def u_vector(Q,parameters):
    cols = Q.shape[1]
    # equivalent to generating a meshgrid
    c1 = [i for i in range(cols) for _ in range(cols)]
    c2 = [i for _ in range(cols) for i in range(cols)]
    return 0.5 * torch.norm(Q[:,c1] - Q[:,c2], dim=0).reshape(cols, cols) * parameters['gamma']

免责声明。使用Q.shape[0]定义u的形状和索引范围是很奇怪的，因为所有索引都用于访问Q的列。上面提供的代码仅在Q是正方形的情况下才与您的函数匹配（另外，如果Q的列多于行，则由于无效索引而使您的函数崩溃）。如果您想要产生与您的输出完全相同的输出的东西，那么下面的方法将起作用。此版本与先前版本之间的唯一区别是，当Q.shape[0] > Q.shape[1]时，它以零填充。

def u_vector(Q,parameters):
    rows, cols = Q.shape[0], Q.shape[1]
    c1 = [i for i in range(cols) for _ in range(cols)]
    c2 = [i for _ in range(cols) for i in range(cols)]
    u = 0.5 * torch.norm(Q[:,c1] - Q[:,c2], dim=0).reshape(cols, cols) * parameters['gamma']
    return torch.nn.functional.pad(u, (0, rows-cols, 0, rows-cols), 'constant', 0)

还要注意，如果Q在GPU上，它将在GPU上运行。

是否有一种向量化将函数应用于张量的每个元素的方法

1 个答案: