| """Load synthetic-pretrained Voronoi microstructure ResNets without pickle.""" |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
| from torchvision import models |
| from safetensors.torch import load_file |
|
|
| TARGET_NAMES = [ |
| 'log_grains', 'mean_ecd', 'area_weighted_ecd', 'ecd_cv', |
| 'hall_petch_mean_ecd', 'phase_area_fraction', 'phase_grain_fraction', |
| 'log_stretch', 'log_defects', 'defect_area_fraction', |
| 'mean_defect_radius', 'boundary_fraction', 'interior_grain_fraction', |
| ] |
| VARIANTS = { |
| 'resnet18_seed1019': 'resnet18', |
| 'resnet18_seed2026': 'resnet18', |
| 'resnet34_seed1019': 'resnet34', |
| 'resnet50_seed1019': 'resnet50', |
| } |
|
|
|
|
| class VoronoiResNet(nn.Module): |
| def __init__(self, architecture): |
| super().__init__() |
| if architecture not in {'resnet18', 'resnet34', 'resnet50'}: |
| raise ValueError('Expected resnet18, resnet34, or resnet50') |
| self.encoder = getattr(models, architecture)(weights=None) |
| dimension = self.encoder.fc.in_features |
| self.encoder.fc = nn.Identity() |
| self.head = nn.Sequential(nn.Linear(dimension, 128), nn.SiLU(), nn.Linear(128, 26)) |
| self.register_buffer('mean', torch.tensor([.485, .456, .406])[None, :, None, None]) |
| self.register_buffer('scale', torch.tensor([.229, .224, .225])[None, :, None, None]) |
| self.register_buffer('target_mean', torch.zeros(13)) |
| self.register_buffer('target_scale', torch.ones(13)) |
|
|
| def forward(self, images): |
| """Return normalized means, clamped log variances, and encoder embeddings.""" |
| if images.ndim != 4 or images.shape[1:] != (1, 256, 256): |
| raise ValueError('Expected floating point images shaped [N, 1, 256, 256]') |
| if not images.is_floating_point(): |
| raise ValueError('Convert image pixels to floating point in [0, 1] first') |
| z = self.encoder((images.repeat(1, 3, 1, 1) - self.mean) / self.scale) |
| p = self.head(z) |
| return p[:, :13], p[:, 13:].clamp(-4, 3), z |
|
|
| @torch.inference_mode() |
| def predict(self, images, tta=True): |
| """Return parameter estimates, standard deviations, and mean embeddings. |
| |
| Log-transformed targets stay in their named log domain. Uncertainty is |
| estimated by the synthetic model and is not calibrated on real images. |
| """ |
| if self.training: |
| raise ValueError('Call model.eval() before inference') |
| means, variances, embeddings = [], [], [] |
| views = ((flip, rotation) for flip in (False, True) for rotation in range(4)) if tta else [(False, 0)] |
| for flip, rotation in views: |
| x = images.flip(-1) if flip else images |
| p, lv, z = self(torch.rot90(x, rotation, (-2, -1))) |
| means.append(p.float()) |
| variances.append(lv.float().exp()) |
| embeddings.append(z.float()) |
| p = torch.stack(means) |
| parameters = p.mean(0) * self.target_scale + self.target_mean |
| variance = torch.stack(variances).mean(0) + p.var(0, unbiased=False) |
| uncertainty = variance.sqrt() * self.target_scale |
| return parameters, uncertainty, torch.stack(embeddings).mean(0) |
|
|
|
|
| def load_model(directory, variant='resnet18_seed1019', device='cpu'): |
| """Load one variant from a downloaded repository directory.""" |
| model = VoronoiResNet(VARIANTS[variant]) |
| model.load_state_dict(load_file(str(Path(directory) / f'{variant}.safetensors')), strict=True) |
| return model.to(device).eval() |
|
|
|
|
| if __name__ == '__main__': |
| torch.set_num_threads(2) |
| model = load_model(Path(__file__).resolve().parent) |
| p, u, z = model.predict(torch.full((1, 1, 256, 256), .6), tta=False) |
| assert p.shape == u.shape == (1, 13) and z.shape == (1, 512) |
| assert torch.isfinite(p).all() and torch.isfinite(u).all() and (u > 0).all() |
| print('Load and inference check passed') |
|
|
