| import json,torch |
| from pathlib import Path |
| from torch import nn |
| import torch.nn.functional as F |
| import yaml |
| def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text()) |
| def sample(i): |
| y,x=torch.meshgrid(torch.linspace(-1,1,96),torch.linspace(-1,1,96),indexing='ij');tc=((x-.3*torch.sin(torch.tensor(i)))**2+(y-.2)**2<.08).long();ar=(abs(y-.4*x)<.08).long()*2;lab=torch.maximum(tc,ar);f=torch.stack((torch.exp(-((x)**2+y**2)),torch.sin(x*4),torch.cos(y*4),(lab>0).float()));return f.float(),lab |
| class ClimateNetDeepLab(nn.Module): |
| def __init__(self,channels=4,classes=3,hidden=16):super().__init__();self.enc=nn.Sequential(nn.Conv2d(channels,hidden,3,padding=1),nn.BatchNorm2d(hidden),nn.ReLU(),nn.Conv2d(hidden,hidden*2,3,2,1),nn.ReLU());self.aspp=nn.ModuleList([nn.Conv2d(hidden*2,hidden,3,padding=d,dilation=d) for d in (1,2,4)]);self.dec=nn.Sequential(nn.Conv2d(hidden*3,hidden,3,padding=1),nn.ReLU(),nn.Conv2d(hidden,classes,1));self.model_config={'channels':channels,'classes':classes,'hidden':hidden} |
| def forward(self,x):z=self.enc(x);z=torch.cat([m(z) for m in self.aspp],1);return F.interpolate(self.dec(z),size=x.shape[-2:],mode='bilinear',align_corners=False) |
| def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n') |
|
|
