时间戳

首页 > 解决方案 > 有没有一种方法可以矢量化将函数应用于张量的每个元素

vectorization - 有没有一种方法可以矢量化将函数应用于张量的每个元素

问题描述

我想在优化过程中实现更新规则。我想使用 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

有人可以帮忙吗？

标签： vectorizationpytorch

解决方案

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

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 上运行。

推荐阅读