| import json,torch,yaml | |
| from pathlib import Path | |
| from torch import nn | |
| def cfg(r):return yaml.safe_load((Path(r)/'conf/config.yaml').read_text()) | |
| class WB2Baseline(nn.Module): | |
| def __init__(self,variables=8,hidden=16):super().__init__();self.n=nn.Sequential(nn.Conv2d(variables,hidden,3,padding=1),nn.ReLU(),nn.Conv2d(hidden,variables,3,padding=1));self.model_config={'variables':variables,'hidden':hidden} | |
| def forward(self,x):return x+self.n(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') | |