| import json,torch |
| from pathlib import Path |
| from torch import nn |
| import yaml |
| AEROSOL=('AOD','TSAOD','SUAOD','DUAOD','BCAOD','OCAOD','SSAOD','SUSMC','DUSMC','BCSMC','OCSMC','SSSMC') |
| def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text()) |
| def state(i,h=0): |
| y,x=torch.meshgrid(torch.linspace(-1,1,16),torch.linspace(-1,1,16),indexing='ij');c=torch.arange(54)[:,None,None];return (torch.sin((c%8+1)*x+.02*(i+h))*torch.cos((c%6+1)*y)).float() |
| class Relay(nn.Module): |
| def __init__(self,c,h,heads):super().__init__();self.e=nn.Conv2d(c,h,2,2);self.a=nn.MultiheadAttention(h,heads,batch_first=True);self.d=nn.ConvTranspose2d(h,c,2,2) |
| def forward(self,x):z=self.e(x);b,c,y,w=z.shape;t=z.permute(0,2,3,1).reshape(b,y*w,c);t=self.a(t,t,t,need_weights=False)[0];return x+self.d(t.reshape(b,y,w,c).permute(0,3,1,2)) |
| class AIGAMFS(nn.Module): |
| def __init__(self,channels=54,hidden=24,heads=4):super().__init__();self.models=nn.ModuleDict({str(h):Relay(channels,hidden,heads) for h in (3,6,9,12)});self.model_config={'channels':channels,'hidden':hidden,'heads':heads} |
| def forward(self,x,lead):return self.models[str(lead)](x) |
| def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n') |
|
|
