| from __future__ import annotations
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
|
|
| class GradReverseFn(torch.autograd.Function):
|
| @staticmethod
|
| def forward(ctx,x,lambd): ctx.lambd=float(lambd); return x.view_as(x)
|
| @staticmethod
|
| def backward(ctx,g): return -ctx.lambd*g,None
|
|
|
| def grad_reverse(x,lambd): return GradReverseFn.apply(x,lambd)
|
|
|
|
|
| class Tower(nn.Module):
|
| def __init__(self,din:int,dout:int,dropout:float=.12):
|
| super().__init__(); self.net=nn.Sequential(nn.LayerNorm(din),nn.Linear(din,dout),nn.GELU(),nn.Dropout(dropout),nn.Linear(dout,dout),nn.GELU())
|
| def forward(self,x): return self.net(x)
|
|
|
|
|
| class RavenFusion(nn.Module):
|
|
|
| BRANCHES=("semantic","camera","spectral","patch","specialist")
|
| def __init__(self,dims:dict,num_ai_domains:int=0,num_real_domains:int=0,latent_dim:int=64):
|
| super().__init__()
|
| outdims={"semantic":96,"camera":48,"spectral":48,"patch":48,"specialist":24}
|
| self.towers=nn.ModuleDict({k:Tower(int(dims[k]),outdims[k]) for k in self.BRANCHES})
|
| self.branch_heads=nn.ModuleDict({k:nn.Linear(outdims[k],1) for k in self.BRANCHES})
|
| total=sum(outdims.values()); self.pre_vib=nn.Sequential(nn.LayerNorm(total),nn.Linear(total,128),nn.GELU(),nn.Dropout(.15))
|
| self.mu=nn.Linear(128,latent_dim); self.logvar=nn.Linear(128,latent_dim)
|
| qdim=int(dims["quality"]); self.router=nn.Sequential(nn.LayerNorm(qdim),nn.Linear(qdim,48),nn.GELU(),nn.Linear(48,len(self.BRANCHES)))
|
| self.anomaly_head=nn.Sequential(nn.LayerNorm(5),nn.Linear(5,24),nn.GELU(),nn.Linear(24,1))
|
| self.global_head=nn.Sequential(nn.LayerNorm(latent_dim),nn.Linear(latent_dim,48),nn.GELU(),nn.Dropout(.12),nn.Linear(48,1))
|
| self.mixer=nn.Linear(3,1)
|
| with torch.no_grad():
|
| self.mixer.weight.copy_(torch.tensor([[1.0,1.0,.8]])); self.mixer.bias.zero_()
|
| self.ai_domain_head=nn.Sequential(nn.Linear(latent_dim,48),nn.GELU(),nn.Linear(48,num_ai_domains)) if num_ai_domains>1 else None
|
| self.real_domain_head=nn.Sequential(nn.Linear(latent_dim,48),nn.GELU(),nn.Linear(48,num_real_domains)) if num_real_domains>1 else None
|
|
|
| def _latent(self,h,sample:bool):
|
| p=self.pre_vib(h); mu=self.mu(p); lv=torch.clamp(self.logvar(p),-6,4)
|
| if sample and self.training: z=mu+torch.randn_like(mu)*torch.exp(.5*lv)
|
| else: z=mu
|
| kl=-.5*torch.mean(1+lv-mu.pow(2)-lv.exp())
|
| return z,mu,lv,kl
|
|
|
| def forward(self,features:dict,quality:torch.Tensor,anchors:torch.Tensor,luma:torch.Tensor,domain_lambda:float=0.0,sample_vib:bool=True):
|
| parts={k:self.towers[k](features[k]) for k in self.BRANCHES}; h=torch.cat([parts[k] for k in self.BRANCHES],1); z,mu,lv,kl=self._latent(h,sample_vib)
|
| branch_logits=torch.cat([self.branch_heads[k](parts[k]) for k in self.BRANCHES],1)
|
|
|
|
|
| route_quality=quality.clone()
|
| if route_quality.shape[1]>=12: route_quality[:,10:12]=0.0
|
| route_logits=self.router(route_quality)
|
|
|
| darkness=torch.clamp((.25-luma)/.25,0,1).unsqueeze(1)
|
| prior=torch.tensor([-.60,.50,.38,.30,.12],dtype=route_logits.dtype,device=route_logits.device).unsqueeze(0)
|
| gates=F.softmax(route_logits+darkness*prior,dim=1); gated=(gates*branch_logits).sum(1,keepdim=True)
|
| global_logit=self.global_head(z); anomaly_logit=self.anomaly_head(anchors)
|
|
|
| components=torch.cat([global_logit,gated,anomaly_logit],1); logit=self.mixer(components).squeeze(1)
|
| ai_domain=self.ai_domain_head(grad_reverse(z,domain_lambda)) if self.ai_domain_head is not None and domain_lambda>0 else None
|
| real_domain=self.real_domain_head(grad_reverse(z,domain_lambda)) if self.real_domain_head is not None and domain_lambda>0 else None
|
| return {"logit":logit,"branch_logits":branch_logits,"gates":gates,"components":components,"latent":mu,"kl":kl,"ai_domain_logits":ai_domain,"real_domain_logits":real_domain}
|
|
|
