这是我的第390篇原创文章。

一、引言

        贝叶斯神经网络的结构与传统神经网络相似，但将权重参数视为随机变量，并赋予先验分布。例如，对于一个简单的用于股票价格预测的多层感知机神经网络，其输出（预测的股票价格）与输入（如历史价格、成交量、宏观经济指标等特征）的关系可以表示为：，其中 是神经网络的前向传播函数， 是权重参数向量，服从某种先验分布，如高斯分布。先验分布的选择依据对股票市场的先验知识或假设，例如认为权重不应过大，以避免过拟合或不稳定的预测。

        在贝叶斯神经网络中，目标是确定神经网络参数的后验分布，而不是像传统神经网络那样确定固定的参数值。在贝叶斯神经网络中，预测不再是一个确定的值，而是一个概率分布。本文实现了训练一个基于贝叶斯神经网络（Bayesian Neural Network，BNN）的小型多层感知机（MLP），并在 MNIST 数据集上进行分类任务。通过使用贝叶斯推断和变分推断方法，模型会进行多次采样，以估计预测的分布，从而为每个样本计算出预测的均值和方差。

二、实现过程

2.1 数据准备

代码：

full_ds = MNIST(data_dir, train=True, download=True, transform=Compose([ToTensor(), Normalize((0.1307,), (0.3081,))]))
train_ds, val_ds = random_split(full_ds, [55000, 5000])
test_ds = MNIST(data_dir, train=False, download=True,transform=Compose([ToTensor(), Normalize((0.1307,), (0.3081,))]))

train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=batch_size, shuffle=False)
test_dl = DataLoader(test_ds, batch_size=batch_size, shuffle=False)

in_dim = train_ds[0][0].numel() # 输入维度
out_dim = 10                    # 输入维度
batch_num=len(train_dl)         # 小批量总数

下载 MNIST 数据集，训练集和验证集的变换操作包括转换为 Tensor 格式，并进行归一化处理。

random_split将完整的训练数据 full_ds 分为 55000 张训练图片和 5000 张验证图片。

下载测试集，并创建训练集、验证集和测试集的 DataLoader。DataLoader 用于将数据分批次传递给模型。

in_dim获取单个输入图像的维度，out_dim是输出的类别数（MNIST 共有10个数字类别）。

batch_num 计算一个 epoch 中的批次总数。

2.2模型定义

代码：

# 模型
model = BayesMLP(in_dim, out_dim, hidden_dims=[128, 64]).to(device)
# 优化器
opt = Adam(model.parameters(), lr=0.001)
# 损失
loss_fn = ClassificationELBOLoss(batch_num=batch_num)

BayesMLP 是一个自定义的贝叶斯多层感知机模型，具有输入维度 in_dim，输出类别数 out_dim，隐藏层大小 [128, 64]。这个模型使用贝叶斯推断进行多次采样。

使用 Adam 优化器来优化模型参数，学习率为 0.001。

ClassificationELBOLoss 是一个自定义损失函数，可能是用于变分推断的损失函数（ELBO：Evidence Lower Bound）。它在每次迭代时根据变分推断的损失进行优化。

2.3 模型训练

代码：

epochs = 3
for epoch in range(epochs):                                                         # 算法2：第2行    
    for batch, (batch_x, batch_y) in enumerate(train_dl):        
        opt.zero_grad()        
        model_out = model(batch_x, sample_num)        
        loss = loss_fn(model_out, batch_y)        
        loss.backward()                                                             # 算法2：第9行        
        opt.step()                                                                  # 算法2：第12行        
        if batch % 100 == 0:            
            acc = evaluate(model, val_dl, sample_num)            
            print(f'epoch: {epoch+1:>2}/{epochs:<2} batch: {batch+1:>4}/{batch_num:<4} Loss: {loss.item():.8}\tValid Acc: {acc: .6}')            
            model.train()

训练循环：训练 3 个 epoch。

batch_x是当前批次的输入数据，batch_y 是标签。

model(batch_x, sample_num)表示使用贝叶斯神经网络模型对输入数据进行预测，并进行 sample_num 次采样

loss_fn(model_out, batch_y)计算模型输出和标签之间的损失。

通过 loss.backward() 和 opt.step() 进行反向传播和参数更新。

2.4 测试可视化

代码：

test_acc = evaluate(model, test_dl, sample_num)
print(f"Test acc: {test_acc:.6}")
# 从测试集中随机选择一个数据
img = random.choice(test_ds)imgs = rotate(img[0], angle=30)      # 旋转图像
model.eval()preds = model(imgs, sample_num=50)                    # 多次采样
preds = torch.stack(preds, dim=-1).softmax(dim=1)
mean_preds = preds.mean(dim=-1).argmax(dim=1)           # 预测的均值
preds_var = preds.var(dim=-1).mean(dim=1)               # 预测的方差
infos = [f'predict: {p.item()}  var: {v.item():.3}' for p, v in zip(mean_preds, preds_var)]grid_show_imgs(imgs, infos=infos)

测试和可视化部分通过旋转图像并多次采样，展示了模型对未知样本的预测分布。

从测试集随机选择一个样本 img。将该图像旋转 30 度。使用模型对旋转后的图像进行预测，进行 50 次采样。计算每个预测的均值（mean_preds）和方差（preds_var）。

构建一个字符串列表 infos，记录每个预测的类别和对应的方差。使用 grid_show_imgs 函数可视化图像和预测信息。

结果：

作者简介：

读研期间发表6篇SCI数据挖掘相关论文，现在某研究院从事数据算法相关科研工作，结合自身科研实践经历不定期分享关于Python、机器学习、深度学习、人工智能系列基础知识与应用案例。致力于只做原创，以最简单的方式理解和学习，关注我一起交流成长。需要数据集和源码的小伙伴可以关注底部公众号添加作者微信。