model/nncam_stable.py

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 ResidualDNN(nn.Module):
 def __init__(self,i,o,h):super().__init__();self.a=nn.Linear(i,h);self.b=nn.Linear(h,h);self.o=nn.Linear(h,o)
 def forward(self,x):z=torch.relu(self.a(x));return self.o(z+torch.relu(self.b(z)))
class NNCAM(nn.Module):
 def __init__(self,hidden=64):super().__init__();self.q=ResidualDNN(122,30,hidden);self.s=ResidualDNN(122,30,hidden);self.rad=ResidualDNN(122,8,hidden);self.model_config={'hidden':hidden}
 def forward(self,x):return torch.cat((self.q(x),self.s(x),self.rad(x)),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')