| from __future__ import annotations |
| from pathlib import Path |
| from typing import Dict |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| @torch.no_grad() |
| def class_moments(residual: torch.Tensor, probs: torch.Tensor, eps: float = 1e-6): |
| """Class-conditional channel moments without materializing [B,C,d,N]. |
| |
| residual: [B,d,H,W,D] |
| probs: [B,C,H,W,D], soft or one-hot class weights at the same resolution. |
| |
| Returns: |
| mu: [C,d] |
| std: [C,d] |
| weights: [C] |
| |
| This implementation uses first and second moments directly and is much more |
| memory efficient than computing (residual - mu)^2 for every class/voxel. |
| """ |
| B, d, H, W, D = residual.shape |
| C = probs.shape[1] |
| r = residual.reshape(B, d, -1) |
| p = probs.reshape(B, C, -1).to(dtype=residual.dtype) |
| weights = p.sum(dim=(0, 2)).clamp_min(eps) |
| mu = torch.einsum("bcn,bdn->cd", p, r) / weights[:, None] |
| second = torch.einsum("bcn,bdn->cd", p, r.pow(2)) / weights[:, None] |
| var = (second - mu.pow(2)).clamp_min(eps) |
| std = torch.sqrt(var) |
| return mu, std, weights |
|
|
|
|
| def hard_onehot(label: torch.Tensor, num_classes: int): |
| return F.one_hot(label.long().clamp(0, num_classes-1), num_classes).permute(0,4,1,2,3).float() |
|
|
|
|
| def save_source_memory(path: str | Path, memory: Dict): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(memory, path) |
|
|
|
|
| def load_source_memory(path: str | Path, device=None): |
| mem = torch.load(path, map_location=device or "cpu") |
| return mem |
|
|
|
|
| def moment_transport_residual(Rt: torch.Tensor, probs: torch.Tensor, source_mu: torch.Tensor, source_std: torch.Tensor, |
| target_mu: torch.Tensor | None = None, target_std: torch.Tensor | None = None, |
| eps: float = 1e-5): |
| """Class-gated diagonal moment transport from target residual stats to source residual stats. |
| |
| Computes R0(u) = sum_c p_c(u) [mu_s_c + sigma_s_c / sigma_t_c * (Rt(u)-mu_t_c)]. |
| |
| The implementation intentionally avoids stacking all class-wise transported |
| residuals, because [B,C,d,H,W,D] can be very large for 3D volumes. |
| """ |
| B, d, H, W, D = Rt.shape |
| C = probs.shape[1] |
| if target_mu is None or target_std is None: |
| target_mu, target_std, _ = class_moments(Rt, probs, eps=eps) |
| source_mu = source_mu.to(Rt.device, Rt.dtype) |
| source_std = source_std.to(Rt.device, Rt.dtype) |
| target_mu = target_mu.to(Rt.device, Rt.dtype) |
| target_std = target_std.to(Rt.device, Rt.dtype) |
| R0 = torch.zeros_like(Rt) |
| for c in range(C): |
| tr = source_mu[c].view(1,d,1,1,1) + (source_std[c].view(1,d,1,1,1) / (target_std[c].view(1,d,1,1,1) + eps)) * (Rt - target_mu[c].view(1,d,1,1,1)) |
| R0 = R0 + probs[:, c:c+1].to(Rt.dtype) * tr |
| return R0, target_mu, target_std |
|
|
