使用 pytorch 的二维矢量化函数

Question

我有一个函数需要在二维火炬张量的每个入口处计算，它取决于两个轴的索引值。现在，我只能将其实现为嵌套的 for 循环，遍历两个轴。这很慢（需要执行 >10^5 次），我想加快速度以更好地缩放。

vs = 200
nt = 12
b = torch.ones(vs)/vs
n_kw = torch.rand((nt, vs))
n_k = torch.rand((nt,))

def estimate_p(nt, vs, n_kw, n_k):
    p = torch.zeros((nt, vs))
    
    for i in range(0, nt):
        for j in range(0, vs):
            p[i,j] = (n_kw[i, j] + b[j])/(n_k[i] + torch.sum(b))
    return p

有没有办法根据 i,j 索引对 this/map 进行矢量化？

Answer 1

试试 broadcasting:

def estimate_p(nt, vs, n_kw, n_k):
    return (n_kw + b) / (n_k + b.sum()).unsqueeze(-1)

使用 pytorch 的二维矢量化函数

2-D Vectorized Function with pytorch

python

vectorization

pytorch

tensor