| import json |
| from pathlib import Path |
| import numpy as np |
| import torch |
| 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 truth(i): |
| y,x=torch.meshgrid(torch.linspace(-1,1,128),torch.linspace(-1,1,128),indexing='ij');return torch.stack([torch.sin((v+1)*x+.1*i)*torch.cos((v+1)*y) for v in range(12)]).float() |
| def observe(x):return F.avg_pool2d(x.reshape(3,4,128,128),16).reshape(12,8,8) |
| class ScoreUNet(nn.Module): |
| def __init__(self,channels=12,hidden=16):super().__init__();self.net=nn.Sequential(nn.Conv2d(channels+1,hidden,3,padding=1),nn.SiLU(),nn.Conv2d(hidden,hidden,3,padding=1),nn.SiLU(),nn.Conv2d(hidden,channels,3,padding=1));self.model_config={'channels':channels,'hidden':hidden} |
| def forward(self,x,sigma):return self.net(torch.cat((x,torch.ones_like(x[:,:1])*sigma),1)) |
| def write(p,o):p=Path(p);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(o,indent=2)+'\n') |
|
|
