import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim import torchvision import torchvision.transforms as transforms from tqdm import tqdm import numpy as np # ------------------------------- # Hyper‑Gas Layer: Gaussian weight distribution # ------------------------------- class HyperGasLinear(nn.Module): """Linear layer where weights and biases are Gaussian distributions. Uses reparameterization trick: w = mu + sigma * epsilon, epsilon ~ N(0,1). """ def __init__(self, in_features, out_features, prior_sigma=0.1): super().__init__() self.in_features = in_features self.out_features = out_features # Learnable mean and variance (log sigma for stability) self.w_mu = nn.Parameter(torch.Tensor(out_features, in_features).normal_(0, 0.02)) self.w_log_sigma = nn.Parameter(torch.Tensor(out_features, in_features).normal_(-3, 0.1)) self.b_mu = nn.Parameter(torch.Tensor(out_features).zero_()) self.b_log_sigma = nn.Parameter(torch.Tensor(out_features).normal_(-3, 0.1)) # Prior distribution (scale mixture of two Gaussians is common, we use fixed sigma prior) self.prior_sigma = prior_sigma def forward(self, x): # Sample weights and biases using reparameterization w_eps = torch.randn_like(self.w_mu) b_eps = torch.randn_like(self.b_mu) w_sigma = torch.exp(self.w_log_sigma) b_sigma = torch.exp(self.b_log_sigma) w = self.w_mu + w_sigma * w_eps b = self.b_mu + b_sigma * b_eps # Compute KL divergence between q(w|theta) and prior p(w) # q = N(mu, sigma^2), p = N(0, prior_sigma^2) # KL = log(prior_sigma/sigma) + (sigma^2 + mu^2)/(2*prior_sigma^2) - 0.5 kl_w = torch.sum(self.w_log_sigma - torch.log(torch.tensor(self.prior_sigma)) + (w_sigma**2 + self.w_mu**2) / (2 * self.prior_sigma**2) - 0.5) kl_b = torch.sum(self.b_log_sigma - torch.log(torch.tensor(self.prior_sigma)) + (b_sigma**2 + self.b_mu**2) / (2 * self.prior_sigma**2) - 0.5) self.kl = kl_w + kl_b return F.linear(x, w, b) class HyperGasConv2d(nn.Module): """2D convolutional layer with Gaussian weights and biases.""" def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, prior_sigma=0.1): super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size self.stride = stride self.padding = padding fan_in = in_channels * self.kernel_size[0] * self.kernel_size[1] self.w_mu = nn.Parameter(torch.Tensor(out_channels, in_channels, *self.kernel_size).normal_(0, 0.02)) self.w_log_sigma = nn.Parameter(torch.Tensor(out_channels, in_channels, *self.kernel_size).normal_(-3, 0.1)) self.b_mu = nn.Parameter(torch.Tensor(out_channels).zero_()) self.b_log_sigma = nn.Parameter(torch.Tensor(out_channels).normal_(-3, 0.1)) self.prior_sigma = prior_sigma def forward(self, x): w_eps = torch.randn_like(self.w_mu) b_eps = torch.randn_like(self.b_mu) w_sigma = torch.exp(self.w_log_sigma) b_sigma = torch.exp(self.b_log_sigma) w = self.w_mu + w_sigma * w_eps b = self.b_mu + b_sigma * b_eps kl_w = torch.sum(self.w_log_sigma - torch.log(torch.tensor(self.prior_sigma)) + (w_sigma**2 + self.w_mu**2) / (2 * self.prior_sigma**2) - 0.5) kl_b = torch.sum(self.b_log_sigma - torch.log(torch.tensor(self.prior_sigma)) + (b_sigma**2 + self.b_mu**2) / (2 * self.prior_sigma**2) - 0.5) self.kl = kl_w + kl_b return F.conv2d(x, w, b, stride=self.stride, padding=self.padding) class HyperGasNet(nn.Module): """Bayesian CNN for CIFAR-10 (32x32).""" def __init__(self, prior_sigma=0.1): super().__init__() self.conv1 = HyperGasConv2d(3, 32, 3, padding=1, prior_sigma=prior_sigma) self.conv2 = HyperGasConv2d(32, 64, 3, padding=1, prior_sigma=prior_sigma) self.pool = nn.MaxPool2d(2, 2) self.fc1 = HyperGasLinear(64 * 8 * 8, 256, prior_sigma=prior_sigma) self.fc2 = HyperGasLinear(256, 10, prior_sigma=prior_sigma) def forward(self, x, sample=True): # If sample==False, use only the mean (deterministic mode, e.g. for eval with single pass) if not sample: # deterministic mode: replace each layer with its mean (mu) only x = F.relu(F.conv2d(x, self.conv1.w_mu, self.conv1.b_mu, stride=1, padding=1)) x = self.pool(x) x = F.relu(F.conv2d(x, self.conv2.w_mu, self.conv2.b_mu, stride=1, padding=1)) x = self.pool(x) x = x.view(x.size(0), -1) x = F.relu(F.linear(x, self.fc1.w_mu, self.fc1.b_mu)) x = F.linear(x, self.fc2.w_mu, self.fc2.b_mu) return x else: # Probabilistic forward: sample weights each time x = F.relu(self.conv1(x)) x = self.pool(x) x = F.relu(self.conv2(x)) x = self.pool(x) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.fc2(x) return x def kl_loss(self): """Sum of KL divergences from all layers.""" kl = (self.conv1.kl + self.conv2.kl + self.fc1.kl + self.fc2.kl) / 60000 # normalize by dataset size return kl # ------------------------------- # Training and evaluation # ------------------------------- def train_epoch(model, device, train_loader, optimizer, epoch, beta=0.1): model.train() total_loss = 0.0 total_correct = 0 for data, target in tqdm(train_loader, desc=f"Epoch {epoch}"): data, target = data.to(device), target.to(device) for _ in range(10): optimizer.zero_grad() # Monte Carlo sample: one forward pass per batch (can be multiple, but time) output = model(data) nll = F.cross_entropy(output, target, reduction='mean') kl = model.kl_loss() loss = nll + beta * kl # beta = 1 / num_batches is common; here we use fixed beta loss.backward() optimizer.step() pred = output.argmax(dim=1, keepdim=True) total_correct += pred.eq(target.view_as(pred)).sum().item() total_loss += loss.item() * len(data) avg_loss = total_loss / len(train_loader.dataset) acc = 100. * total_correct / len(train_loader.dataset) print(f"Train Epoch {epoch}: Loss={avg_loss:.4f}, Accuracy={acc:.2f}%") return avg_loss, acc def test(model, device, test_loader, num_samples=10): """Predict using multiple Monte Carlo samples (gas phase aggregation).""" model.eval() correct = 0 with torch.no_grad(): for data, target in tqdm(test_loader, desc="Testing"): data, target = data.to(device), target.to(device) # Average predictions over several weight samples outputs = torch.zeros(data.size(0), 10).to(device) for _ in range(num_samples): outputs += F.softmax(model(data), dim=1) outputs /= num_samples pred = outputs.argmax(dim=1) correct += pred.eq(target).sum().item() acc = 100. * correct / len(test_loader.dataset) print(f"Test accuracy (MC {num_samples} samples): {acc:.2f}%") return acc def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) train_set = torchvision.datasets.CIFAR10(root='../data', train=True, download=True, transform=transform_train) test_set = torchvision.datasets.CIFAR10(root='../data', train=False, download=True, transform=transform_test) train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2) test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False, num_workers=2) model = HyperGasNet(prior_sigma=0.1).to(device) optimizer = optim.Adam(model.parameters(), lr=1e-3) for epoch in range(1, 31): train_epoch(model, device, train_loader, optimizer, epoch, beta=0.1) if epoch % 5 == 0: test(model, device, test_loader, num_samples=10) # Final evaluation with many samples (gas phase) print("\nFinal evaluation with 50 Monte Carlo samples:") test(model, device, test_loader, num_samples=50) if __name__ == "__main__": main()