pytorch - pytorch中的内存高效批量成对KL散度

Question

给定一个 Batch 张量S维度BxD，其中每一行都是维度上的概率分布D，我想计算成对的 KL 散度矩阵KL（维度BxB），使得KL[i,j] = KL-DIVERGENCE(S[i],S[j]).

这有效：

softmax = torch.nn.Softmax(dim = 1)
B,D = 100, 1000
S = softmax(torch.rand(B,D))
KL = torch.mean(S[:,None,:] * (S[:,None,:]/S).log(),dim = 2)

但是内存效率很低。有没有办法更有效地计算这个矩阵？

谢谢。

标签： pytorchtensorinformation-theory