Pytorch实现一个简单的贝叶斯卷积神经网络模型

贝叶斯深度模型的主要特点和实现说明：

1. 模型结构：

   • 结合了常规卷积层（用于特征提取）和贝叶斯线性层（用于分类）
   • 贝叶斯层将权重视为随机变量，而非传统神经网络中的确定值
   • 使用变分推断来近似权重的后验分布

2. 贝叶斯特性：

   • 通过重参数化技巧实现随机变量的采样，使得模型可训练
   • 损失函数包含两部分：分类损失（交叉熵）和 KL 散度（衡量近似后验与先验的差异）
   • 测试时通过多次采样获取预测分布，体现模型的不确定性

3. 使用方法：

   • 代码会自动下载 MNIST 数据集并进行预处理
   • 支持 GPU 加速（如果可用）
   • 训练完成后会绘制损失和准确率曲线，并保存模型

4. 与传统神经网络的区别：

   • 贝叶斯模型能够提供预测的不确定性估计
   • 通常具有更好的泛化能力，不易过拟合
   • 训练过程更复杂，计算成本更高

import torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transformsimport numpy as npimport matplotlib.pyplot as plt# 定义贝叶斯线性层 - 使用变分推断近似后验分布class BayesianLinear(nn.Module): def __init__(self, in_features, out_features): super(BayesianLinear, self).__init__() self.in_features = in_features self.out_features = out_features # 先验分布参数 (高斯分布) self.prior_mu = 0.0 self.prior_sigma = 1.0 # 变分参数 - 权重的均值和标准差 self.mu_weight = nn.Parameter(torch.Tensor(out_features, in_features).normal_(0, 0.1)) self.sigma_weight = nn.Parameter(torch.Tensor(out_features, in_features).fill_(0.1)) # 变分参数 - 偏置的均值和标准差 self.mu_bias = nn.Parameter(torch.Tensor(out_features).normal_(0, 0.1)) self.sigma_bias = nn.Parameter(torch.Tensor(out_features).fill_(0.1)) # 用于重参数化技巧的噪声变量 self.epsilon_weight = None self.epsilon_bias = None def forward(self, x): # 重参数化技巧：将随机采样转换为确定性操作，便于反向传播 if self.training: # 训练时从近似后验分布中采样 self.epsilon_weight = torch.normal(torch.zeros_like(self.mu_weight)) self.epsilon_bias = torch.normal(torch.zeros_like(self.mu_bias)) weight = self.mu_weight + self.sigma_weight * self.epsilon_weight bias = self.mu_bias + self.sigma_bias * self.epsilon_bias else: # 测试时使用均值（最大后验估计） weight = self.mu_weight bias = self.mu_bias # 计算KL散度（衡量近似后验与先验的差异） kl_loss = self._kl_divergence() return nn.functional.linear(x, weight, bias), kl_loss def _kl_divergence(self): # 计算KL散度：KL(q(w) || p(w)) kl_weight = 0.5 * torch.sum( 1 + 2 * torch.log(self.sigma_weight) - torch.square(self.mu_weight) - torch.square(self.sigma_weight) ) / (self.prior_sigma ** 2) kl_bias = 0.5 * torch.sum( 1 + 2 * torch.log(self.sigma_bias) - torch.square(self.mu_bias) - torch.square(self.sigma_bias) ) / (self.prior_sigma ** 2) return kl_weight + kl_bias# 定义贝叶斯卷积神经网络模型class BayesianCNN(nn.Module): def __init__(self, num_classes=10): super(BayesianCNN, self).__init__() # 卷积层使用常规卷积（为简化模型） self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # 全连接层使用贝叶斯层 self.fc1 = BayesianLinear(64 * 7 * 7, 128) self.fc2 = BayesianLinear(128, num_classes) self.relu = nn.ReLU() def forward(self, x): # 卷积特征提取部分 x = self.pool(self.relu(self.conv1(x))) x = self.pool(self.relu(self.conv2(x))) x = x.view(-1, 64 * 7 * 7) # 展平特征图 # 贝叶斯全连接部分 x, kl1 = self.fc1(x) x = self.relu(x) x, kl2 = self.fc2(x) # 总KL散度 total_kl = kl1 + kl2 return x, total_kl# 训练函数def train(model, train_loader, optimizer, criterion, epoch, device): model.train() train_loss = 0 correct = 0 total = 0 # KL散度的权重（根据数据集大小调整） kl_weight = 1.0 / len(train_loader.dataset) for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() # 前向传播 output, kl_loss = model(data) # 总损失 = 分类损失 + KL散度正则化 loss = criterion(output, target) + kl_weight * kl_loss # 反向传播和优化 loss.backward() optimizer.step() # 统计 train_loss += loss.item() _, predicted = torch.max(output.data, 1) total += target.size(0) correct += (predicted == target).sum().item() # 打印训练进度 if batch_idx % 100 == 0: print(f\'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} \'  f\'({100. * batch_idx / len(train_loader):.0f}%)]\\tLoss: {loss.item():.6f}\') train_loss /= len(train_loader) train_acc = 100. * correct / total print(f\'Train set: Average loss: {train_loss:.4f}, Accuracy: {correct}/{total} ({train_acc:.2f}%)\') return train_loss, train_acc# 测试函数def test(model, test_loader, criterion, device, num_samples=10): model.eval() test_loss = 0 correct = 0 total = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) # 多次采样以获取预测分布（体现贝叶斯模型的不确定性） outputs = [] for _ in range(num_samples): output, _ = model(data) outputs.append(output.unsqueeze(0)) # 平均多次采样的结果 output = torch.mean(torch.cat(outputs, dim=0), dim=0) test_loss += criterion(output, target).item() # 统计准确率 _, predicted = torch.max(output.data, 1) total += target.size(0) correct += (predicted == target).sum().item() test_loss /= len(test_loader) test_acc = 100. * correct / total print(f\'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{total} ({test_acc:.2f}%)\') return test_loss, test_acc# 主函数def main(): # 超参数设置 batch_size = 64 test_batch_size = 1000 epochs = 10 lr = 0.001 seed = 42 num_samples = 10 # 测试时的采样次数，用于获取预测分布 # 设置设备（GPU或CPU） device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\") print(f\"Using device: {device}\") # 设置随机种子，保证结果可复现 torch.manual_seed(seed) # 数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) # MNIST数据集的均值和标准差 ]) # 加载MNIST数据集 train_dataset = datasets.MNIST( root=\'./data\', train=True, download=True, transform=transform ) test_dataset = datasets.MNIST( root=\'./data\', train=False, download=True, transform=transform ) # 创建数据加载器 train_loader = DataLoader( train_dataset, batch_size=batch_size, shuffle=True ) test_loader = DataLoader( test_dataset, batch_size=test_batch_size, shuffle=False ) # 初始化模型、损失函数和优化器 model = BayesianCNN().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=lr) # 记录训练过程中的损失和准确率 train_losses = [] train_accs = [] test_losses = [] test_accs = [] # 开始训练和测试 for epoch in range(1, epochs + 1): train_loss, train_acc = train(model, train_loader, optimizer, criterion, epoch, device) test_loss, test_acc = test(model, test_loader, criterion, device, num_samples) train_losses.append(train_loss) train_accs.append(train_acc) test_losses.append(test_loss) test_accs.append(test_acc) # 绘制训练和测试损失曲线 plt.figure(figsize=(12, 5)) plt.subplot(1, 2, 1) plt.plot(range(1, epochs + 1), train_losses, label=\'Train Loss\') plt.plot(range(1, epochs + 1), test_losses, label=\'Test Loss\') plt.xlabel(\'Epoch\') plt.ylabel(\'Loss\') plt.title(\'Loss vs Epoch\') plt.legend() # 绘制训练和测试准确率曲线 plt.subplot(1, 2, 2) plt.plot(range(1, epochs + 1), train_accs, label=\'Train Accuracy\') plt.plot(range(1, epochs + 1), test_accs, label=\'Test Accuracy\') plt.xlabel(\'Epoch\') plt.ylabel(\'Accuracy (%)\') plt.title(\'Accuracy vs Epoch\') plt.legend() plt.tight_layout() plt.show() # 保存模型 torch.save(model.state_dict(), \'bayesian_cnn_mnist.pth\') print(\"Model saved as \'bayesian_cnn_mnist.pth\'\")if __name__ == \'__main__\': main()

在模型规模相似（例如参数总量、网络深度和宽度相近）的情况下，普通卷积神经网络（CNN）的训练效率通常更高，训练速度更快。这主要源于贝叶斯卷积神经网络（Bayesian CNN）的特殊结构和训练机制带来的额外计算开销，具体原因如下：

1. 参数数量与计算复杂度差异

普通 CNN 中，每个权重是确定值，每个层仅需存储和优化一组权重参数（例如卷积核权重、偏置）。
而贝叶斯 CNN 中，权重被视为随机变量（通常假设服从高斯分布），需要用变分推断近似其 posterior 分布。这意味着每个权重需要学习两个参数：均值（μ） 和标准差（σ）（或精度），参数数量几乎是普通 CNN 的 2 倍（对于贝叶斯层而言）。

更多的参数直接导致：

• 前向传播时需要计算更多变量的组合（例如通过重参数化技巧采样权重：weight = μ + σ·ε）；
• 反向传播时需要计算更多参数的梯度（不仅是均值，还有标准差），增加了梯度计算的复杂度。

2. 额外的损失项计算

普通 CNN 的损失函数通常仅包含任务相关损失（例如分类问题的交叉熵损失）。
而贝叶斯 CNN 的损失函数必须包含两部分：

• 任务相关损失（与普通 CNN 相同）；
• KL 散度（KL divergence）：用于衡量近似后验分布与先验分布的差异，作为正则化项。

KL 散度的计算需要对每个贝叶斯层的权重分布进行积分近似（即使是简化的解析解，也需要对所有权重的均值和标准差进行逐元素运算），这会额外增加计算开销，尤其当贝叶斯层较多时，累积开销显著。

3. 采样操作的开销

贝叶斯 CNN 在训练时，为了通过重参数化技巧实现梯度回传，需要对每个贝叶斯层的权重进行随机采样（例如从N(μ, σ²)中采样噪声ε，再计算weight = μ + σ·ε）。虽然采样操作本身不算复杂，但在大规模网络中，多次采样（即使每个 batch 一次）会累积计算时间。

普通 CNN 则无需采样，权重是确定性的，前向传播更直接高效。

总结

在模型规模相似的情况下，普通 CNN 由于参数更少、计算流程更简单（无额外的 KL 散度计算和采样操作），训练速度显著快于贝叶斯 CNN。

贝叶斯 CNN 的优势不在于训练效率，而在于其能量化预测的不确定性（例如通过多次采样得到预测分布），并在小样本、数据噪声大的场景下可能具有更好的泛化能力，但这是以更高的计算成本为代价的。