Publish ClimateNet reproduction

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pt binary
+*.npz binary

conf/config.yaml

@@ -0,0 +1,6 @@
+seed: 42
+data: {path: data/climatenet.npz, samples: 24, channels: 4, tile: [96, 96], logical_grid: [1152, 768], classes: 3}
+model: {channels: 4, classes: 3, hidden: 16}
+train: {epochs: 2, learning_rate: 0.0015, class_weights: [1, 16, 16]}
+paths: {checkpoint: result/checkpoints/climatenet.pt, training_metrics: result/training/metrics.json, predictions: result/output/predictions.npz, evaluation: result/evaluation/metrics.json, figure: result/evaluation/comparison.png}
+paper_model: {input_shape: [4, 1152, 768], classes: [BG, TC, AR], optimizer: Adam, learning_rate: 0.0015, batch_size: 16, expert_epochs: 5, paper_license: CC-BY-4.0}

config.json

@@ -0,0 +1 @@
+{"model_name":"ClimateNet","model_type":"climatenet","architectures":["ClimateNetDeepLab"],"framework":"PyTorch","domain":"extreme-weather","task":"semantic-segmentation","implementation":{"entry_point":"model/climatenet.py","scope":"L2 DeepLabv3+ tile reproduction"},"architecture":{"input_channels":["TMQ","U850","V850","PRECT"],"classes":["background","TC","AR"],"logical_grid":[1152,768],"core":"DeepLabv3+ encoder-decoder"},"data":{"source":"CAM5.1","resolution_km":25,"expert_samples":459,"split":[422,18,19],"synthetic_tiles":true},"configuration_sources":["conf/config.yaml","model/climatenet.py","scripts/fake_data.py","scripts/train.py","scripts/inference.py","scripts/result.py"]}

model/climatenet.py

@@ -0,0 +1,12 @@
+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')

scripts/fake_data.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+import sys,numpy as np
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.climatenet import *
+c=cfg(R);x=[];y=[]
+for i in range(24):a,b=sample(i);x.append(a);y.append(b)
+p=R/c['data']['path'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,input=x,target=y,split=[0]*16+[1]*4+[2]*4,origins=np.array([[i*20%1056,i*25%672] for i in range(24)]),logical_shape=[4,1152,768],is_complete_global=False);print(p)

scripts/inference.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+import sys,numpy as np,torch
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.climatenet import *
+c=cfg(R);d=np.load(R/c['data']['path']);z=torch.load(R/c['paths']['checkpoint'],map_location='cpu',weights_only=True);m=ClimateNetDeepLab(**z['model_config']);m.load_state_dict(z['model']);mask=d['split']==2
+with torch.no_grad():logits=m(torch.tensor(d['input'][mask]));pred=logits.argmax(1)
+p=R/c['paths']['predictions'];p.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(p,prediction=pred.numpy(),probability=logits.softmax(1).numpy(),target=d['target'][mask],origins=d['origins'][mask],logical_shape=d['logical_shape'],is_complete_global=False);print(p)

scripts/result.py

@@ -0,0 +1,6 @@
+from pathlib import Path
+import sys,numpy as np;import matplotlib;matplotlib.use('Agg');import matplotlib.pyplot as plt
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.climatenet import cfg,write
+c=cfg(R);d=np.load(R/c['paths']['predictions']);p=d['prediction'];t=d['target'];ious=[]
+for k in range(3):inter=((p==k)&(t==k)).sum();union=((p==k)|(t==k)).sum();ious.append(float(inter/max(union,1)))
+write(R/c['paths']['evaluation'],{'class_iou':ious,'mean_iou':float(np.mean(ious)),'is_complete_global':False,'synthetic':True});fig,ax=plt.subplots(1,2,figsize=(9,4));ax[0].imshow(t[0]);ax[1].imshow(p[0]);q=R/c['paths']['figure'];q.parent.mkdir(parents=True,exist_ok=True);plt.savefig(q,dpi=150);print(q)

scripts/train.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+import sys,os,numpy as np,torch;import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+R=Path(__file__).resolve().parents[1];sys.path.insert(0,str(R));from model.climatenet import *
+c=cfg(R);rank=int(os.getenv('RANK',0));world=int(os.getenv('WORLD_SIZE',1));ddp=world>1
+if ddp:dist.init_process_group('gloo')
+d=np.load(R/c['data']['path']);ids=np.where(d['split']==0)[0];base=ClimateNetDeepLab(**c['model']);m=DDP(base) if ddp else base;opt=torch.optim.Adam(m.parameters(),lr=c['train']['learning_rate']);w=torch.tensor(c['train']['class_weights']).float();ls=[]
+for _ in range(c['train']['epochs']):
+ for i in ids[rank::world]:loss=F.cross_entropy(m(torch.tensor(d['input'][i:i+1])),torch.tensor(d['target'][i:i+1]),weight=w);opt.zero_grad();loss.backward();opt.step();ls.append(float(loss))
+v=torch.tensor([sum(ls),len(ls)],dtype=torch.float64)
+if ddp:dist.all_reduce(v)
+p=R/c['paths']['checkpoint']
+if rank==0:p.parent.mkdir(parents=True,exist_ok=True);torch.save({'model':base.state_dict(),'model_config':c['model']},p);write(R/c['paths']['training_metrics'],{'weighted_cross_entropy':float(v[0]/v[1]),'world_size':world});print(p)
+if ddp:dist.destroy_process_group()