python - 在 Pytorch 中按元素应用 Kullback-Leibler(又名 kl 散度)

Question

我有两个名为 x_t、x_k 的张量，分别具有以下形状 NxHxW 和 KxNxHxW，其中 K，是用于重建x_t的自动编码器的数量(如果您不知道这是什么，假设它们是K个不同的网络，旨在预测x_t，这可能与问题无关)N 是批量大小，H 矩阵高度， W 矩阵宽度。

我正在尝试申请Kullback-Leibler divergence使用 Pytorch's nn.functional.kl_div 对两个张量进行算法(在沿 K^th 维度将 x_t 广播为 x_k 后)方法。

但是，它似乎没有按我的预期工作。。我正在计算 x_t 和 x_k 中每个观察值之间的 kl_div ，从而得到大小为 KxN 的张量(即每个K自动编码器的每个观察的kl_div)。

如果我使用reduction参数，实际输出是一个单个值，如果我使用相同的张量大小(即KxNxHxW)不要使用它。

有人尝试过类似的方法吗？

---

可重现的示例:

import torch
import torch.nn.functional as F
#                  K   N   H  W
x_t = torch.randn(    10, 5, 5)
x_k = torch.randn( 3, 10, 5, 5)

x_broadcasted = x_t.expand_as(x_k)

loss = F.kl_div(x_t, x_k, reduction="none") # or "batchmean", or there are many options

Answer 1

我不清楚你的模型中的概率分布到底是什么。使用 reduction='none'、kl_div，给定 log(x_n) 和 y_n，计算 kl_div = y_n * (log(y_n) - log(x_n))，这是实际 Kullback-Leibler 散度的“求和”部分。求和(或者换句话说，取期望)取决于您。如果您的观点是 H、W 是您想要获取期望值的两个维度，那么就很简单

loss = F.kl_div(x_t, x_k, reduction="none").sum(dim=(-1, -2))

其形状为[K, N]。如果要以不同方式解释网络输出，则需要更好地指定哪些是事件维度，哪些是分布的样本维度。