PyTorch自定义NLL损失仅在batch_size=1时生效问题排查

我仔细检查了你的自定义NLL损失函数，发现了两个核心问题导致大batch_size下模型无法拟合：

1. 对数方差先验项的比例错误 #

在贝叶斯神经网络的ELBO（证据下界）计算中，当使用mini-batch训练时，每个batch的损失需要对应整个训练集的比例。你的代码中，对数方差先验项的计算只取了当前batch的均值，然后直接除以总样本数，这相当于把先验项的贡献缩小了batch_size倍。

举个例子：当batch_size=20时，log_variance_prior返回的是20个样本的先验均值，除以总样本数后得到的是(sum_prior / 20) / N，但正确的计算应该是sum_prior / N，也就是需要乘以batch_size来还原总和。

2. 潜在的设备不匹配问题 #

你的损失函数中，创建的torch.tensor(mean)和torch.tensor(variance)默认在CPU上，如果模型部署在GPU上，会导致张量设备不匹配的错误，这虽然不影响CPU训练，但会限制代码的扩展性。

修改后的NLLLoss代码 #

class NLLLoss(torch.nn.modules.loss._Loss):
    def __init__(self, parameters, num_datapoints, size_average=False, reduce=True):
        super().__init__(size_average, reduce)
        self.parameters = tuple(parameters)
        self.num_datapoints = num_datapoints

    def log_variance_prior(self, log_variance, mean=1e-6, variance=0.01):
        # 将先验参数张量移到输入同设备
        mean_tensor = torch.tensor(mean, device=log_variance.device)
        variance_tensor = torch.tensor(variance, device=log_variance.device)
        # 去掉不必要的dim=1 sum，因为每个样本的log_variance是1维的
        return torch.mean(
            ((-(log_variance - torch.log(mean_tensor)) ** 2) / (2. * variance_tensor)) 
            - 0.5 * torch.log(variance_tensor)
        )

    def weight_prior(self, parameters, wdecay=1.):
        num_parameters = 0
        log_likelihood = 0.
        for parameter in parameters:
            param_count = torch.prod(torch.tensor(parameter.size(), device=parameter.device))
            num_parameters += param_count
            log_likelihood += torch.sum(-wdecay * 0.5 * (parameter ** 2))
        # 避免创建不必要的张量列表，直接累加更高效
        return log_likelihood / (num_parameters.float() + 1e-16)

    def forward(self, input, target):
        torch.nn.modules.loss._assert_no_grad(target)
        batch_size, *_ = input.shape
        prediction_mean = input[:, 0].view(-1, 1)
        log_prediction_variance = input[:, 1].view(-1, 1)
        prediction_variance_inverse = 1. / (torch.exp(log_prediction_variance) + 1e-16)
        mean_squared_error = torch.pow(target - prediction_mean, 2)
        
        # 简化log似然计算：直接对batch求和，再除以batch_size得到平均
        log_likelihood = torch.sum(
            -mean_squared_error * 0.5 * prediction_variance_inverse 
            - 0.5 * log_prediction_variance
        )
        log_likelihood /= batch_size
        
        # 修正先验项比例：乘以batch_size还原总和，再除以总样本数
        log_likelihood += (
            self.log_variance_prior(log_prediction_variance) * batch_size 
            / self.num_datapoints
        )
        log_likelihood += self.weight_prior(self.parameters) / self.num_datapoints
        
        return -log_likelihood

为什么batch_size=1时能正常工作？ #

当batch_size=1时，log_variance_prior返回的是单个样本的先验值，乘以1后除以总样本数，正好和原Keras代码的计算逻辑一致，所以结果正确。而当batch_size增大时，比例错误导致先验项的贡献被严重缩小，模型无法正确学习到方差和权重的正则化约束，最终拟合失败。

额外优化建议 #

• 你的AppendLayer可以简化，直接用nn.Parameter生成方差偏置，避免多余的张量拼接操作；
• 训练时可以加入梯度裁剪，防止大batch下梯度爆炸；
• PyTorch新版本中_assert_no_grad已被标记为过时，建议改用torch._assert或者直接忽略（PyTorch会自动处理梯度传播）。

内容的提问来源于stack exchange，提问作者Nablezen