File size: 1,828 Bytes
dd39c71 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | import torch
import torch.nn.functional as F
import json
import os
from huggingface_hub import hf_hub_download
from model import PhysicsGuidedBNN
class PGBNNPredictor:
"""
Inference and Uncertainty Quantification engine.
"""
def __init__(self, model: PhysicsGuidedBNN, config: dict):
self.model = model
self.config = config
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.device)
self.model.eval()
@classmethod
def from_pretrained(cls, repo_id: str):
config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
with open(config_path, "r") as f:
config = json.load(f)
model = PhysicsGuidedBNN(config)
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
return cls(model, config)
def predict(self, x: torch.Tensor, num_samples: int = 30):
"""
Performs Monte Carlo forward passes to compute predictive mean and epistemic uncertainty.
"""
x = x.to(self.device)
probs_list = []
with torch.no_grad():
for _ in range(num_samples):
logits = self.model(x)
probs = F.softmax(logits, dim=-1)[:, 1] # Probability of fault (class 1)
probs_list.append(probs.unsqueeze(0))
probs_tensor = torch.cat(probs_list, dim=0) # [num_samples, batch_size]
# Mean prediction across MC samples
mean_prob = torch.mean(probs_tensor, dim=0)
# Epistemic Uncertainty (Standard Deviation across MC samples)
epistemic_uncertainty = torch.std(probs_tensor, dim=0)
return mean_prob, epistemic_uncertainty |