如何优化多变量样本伯努利对数似然的PyTorch向量化计算

完全不需要任何for循环，即可实现纯向量化的伯努利对数似然计算，运算效率远高于循环实现。

实现思路 #

• 两个张量p和x形状均为(batch_size, input_size)，所有运算均为逐元素运算，可直接利用PyTorch的广播机制完成，无需遍历批次或特征维度
• 注意要对p做数值裁剪，避免p=0或p=1时计算对数出现inf/nan问题

代码实现 #

手动实现版本 #

import torch

def bernoulli_log_likelihood(p: torch.Tensor, x: torch.Tensor, eps: float = 1e-10) -> torch.Tensor:
    # 裁剪p的数值范围，避免对数计算异常
    p_clamped = torch.clamp(p, min=eps, max=1 - eps)
    # 逐元素计算对数似然项
    log_likelihood_item = x * torch.log(p_clamped) + (1 - x) * torch.log(1 - p_clamped)
    # 按特征维度求和，返回形状为(batch_size)的张量
    return torch.sum(log_likelihood_item, dim=1)

内置API实现版本（数值稳定性更优） #

PyTorch内置的二元交叉熵接口本质就是负的伯努利对数似然，只需调整参数即可直接调用：

import torch.nn.functional as F

def bernoulli_log_likelihood_builtin(p: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
    # reduction='none'保留所有逐元素结果，求和后取负即可得到对数似然
    return -torch.sum(F.binary_cross_entropy(p, x, reduction='none'), dim=1)

结果验证 #

用你给出的示例输入测试：

p = torch.tensor([[0.6, 0.4, 0.], [0.33, 0.34, 0.33]])
x = torch.tensor([[1., 1., 0.], [0., 1., 1.]])
print(bernoulli_log_likelihood(p, x))

输出结果和你预期的计算结果完全一致。

内容的提问来源于stack exchange，提问作者R.DM