| """Compact residual-inception squeeze-excitation UNet.""" |
| import json,math |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| import yaml |
|
|
| def load_config(root):return yaml.safe_load((Path(root)/"conf/config.yaml").read_text()) |
| def synthetic_initialization(i,members=11): |
| y,x=torch.meshgrid(torch.linspace(-1,1,48),torch.linspace(-2,2,96),indexing="ij");base=torch.sin(2*x)*torch.cos(2*y)+.2*i;xs=[] |
| for m in range(members):xs.append(torch.stack([base+.03*m+.05*j+torch.sin(x*(j+1))*.1 for j in range(8)])) |
| targets=torch.stack([base-.12*w+.08*torch.sin(x*w+y) for w in range(1,6)]);return torch.stack(xs),targets |
|
|
| class RISEBlock(nn.Module): |
| def __init__(self,cin,cout,dropout): |
| super().__init__();q=max(1,cout//4);self.paths=nn.ModuleList([nn.Conv2d(cin,q,k,padding=k//2) for k in (3,5,7)]);self.pool=nn.Sequential(nn.MaxPool2d(5,1,2),nn.Conv2d(cin,cout-3*q,1));self.norm=nn.BatchNorm2d(cout);self.skip=nn.Conv2d(cin,cout,1);self.se=nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(cout,max(1,cout//4),1),nn.ReLU(),nn.Conv2d(max(1,cout//4),cout,1),nn.Sigmoid());self.drop=nn.Dropout2d(dropout) |
| def forward(self,x): |
| z=torch.cat([p(x) for p in self.paths]+[self.pool(x)],1);z=F.relu(self.norm(z));z=self.drop(z);return F.relu(z+self.skip(x))*self.se(z) |
| class RISEUNet(nn.Module): |
| def __init__(self,input_channels=8,width=12,dropout=.15): |
| super().__init__();self.e1=RISEBlock(input_channels,width,dropout);self.e2=RISEBlock(width,width*2,dropout);self.e3=RISEBlock(width*2,width*4,dropout);self.d2=RISEBlock(width*6,width*2,dropout);self.d1=RISEBlock(width*3,width,dropout);self.deep=nn.Conv2d(width*2,1,1);self.out=nn.Conv2d(width,1,1);self.model_config={"input_channels":input_channels,"width":width,"dropout":dropout} |
| def forward(self,x): |
| if x.shape[1:]!=(8,48,96):raise ValueError("expected [B,8,48,96]") |
| a=self.e1(x);b=self.e2(F.max_pool2d(a,2));c=self.e3(F.max_pool2d(b,2));u=self.d2(torch.cat((F.interpolate(c,size=b.shape[-2:],mode="bilinear"),b),1));v=self.d1(torch.cat((F.interpolate(u,size=a.shape[-2:],mode="bilinear"),a),1));d=F.interpolate(self.deep(u),size=(48,96),mode="bilinear");o=self.out(v);return (d,o,(d+o)/2) |
| def crps_exp(outputs,target,members=11,f=.08): |
| pred=outputs[-1].reshape(-1,members,48,96);truth=target.reshape(-1,members,48,96);return (abs(pred-truth).mean(1)-f*pred.std(1)).mean()+sum(abs(o-target).mean() for o in outputs[:-1])*.1 |
| def write_json(path,obj):path=Path(path);path.parent.mkdir(parents=True,exist_ok=True);path.write_text(json.dumps(obj,indent=2)+"\n") |
|
|
