带有 pytorch 的二维矢量化函数

问题描述

我有一个函数，我需要在二维火炬张量的每个条目处计算它，它取决于两个轴的索引值。现在，我只能将其实现为嵌套的 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 索引对这个/地图进行矢量化处理？

解决方法

尝试使用 broadcasting：

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