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

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

1. 模型结构：

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

2. 贝叶斯特性：

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

3. 使用方法：

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

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

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

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
import 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 的优势不在于训练效率，而在于其能量化预测的不确定性（例如通过多次采样得到预测分布），并在小样本、数据噪声大的场景下可能具有更好的泛化能力，但这是以更高的计算成本为代价的。