| import torch |
| import torch.nn as nn |
| import numpy as np |
| from pathlib import Path |
| from typing import Optional |
|
|
|
|
| class BetaRegressor(nn.Module): |
| def __init__(self, input_dim: int = 9, output_dim: int = 10, hidden_dims: list = [64, 32]): |
| super(BetaRegressor, self).__init__() |
| |
| layers = [] |
| prev_dim = input_dim |
| |
| for hidden_dim in hidden_dims: |
| layers.append(nn.Linear(prev_dim, hidden_dim)) |
| layers.append(nn.ReLU()) |
| layers.append(nn.Dropout(0.1)) |
| prev_dim = hidden_dim |
| |
| layers.append(nn.Linear(prev_dim, output_dim)) |
| layers.append(nn.Tanh()) |
| |
| self.network = nn.Sequential(*layers) |
| self._initialize_weights() |
| |
| def _initialize_weights(self): |
| for m in self.modules(): |
| if isinstance(m, nn.Linear): |
| nn.init.xavier_uniform_(m.weight, gain=0.1) |
| nn.init.zeros_(m.bias) |
| |
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.network(x) |
|
|
|
|
| class MeasurementToBetaPredictor: |
| def __init__(self, model_path: Optional[str] = None, device: str = "cpu"): |
| self.device = torch.device(device) |
| self.model = BetaRegressor().to(self.device) |
| self.model.eval() |
| |
| if model_path and Path(model_path).exists(): |
| self.load_model(model_path) |
| else: |
| print("Warning: Using untrained model. Results may not be optimal.") |
| print("Consider training the model or loading pretrained weights.") |
| |
| def load_model(self, model_path: str): |
| checkpoint = torch.load(model_path, map_location=self.device) |
| if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint: |
| self.model.load_state_dict(checkpoint['model_state_dict']) |
| else: |
| self.model.load_state_dict(checkpoint) |
| print(f"Loaded model from {model_path}") |
| |
| def predict(self, normalized_measurements: np.ndarray) -> np.ndarray: |
| with torch.no_grad(): |
| measurements_tensor = torch.FloatTensor(normalized_measurements).unsqueeze(0).to(self.device) |
| betas_tensor = self.model(measurements_tensor) |
| betas = betas_tensor.squeeze(0).cpu().numpy() |
| betas = betas * 2.0 |
| |
| return betas |
|
|
|
|
| _predictor_instance = None |
|
|
|
|
| def get_predictor(model_path: Optional[str] = None, device: str = "cpu") -> MeasurementToBetaPredictor: |
| global _predictor_instance |
| if _predictor_instance is None: |
| _predictor_instance = MeasurementToBetaPredictor(model_path=model_path, device=device) |
| return _predictor_instance |
|
|
|
|
| def predict_betas(normalized_measurements: np.ndarray, model_path: Optional[str] = None) -> np.ndarray: |
| predictor = get_predictor(model_path=model_path) |
| return predictor.predict(normalized_measurements) |
|
|
|
|
