u52f85420
04bf90c09b
- 新增 irstd 核心包、scripts、tests、configs - 新增研究文档：literature_review / mathematical_analysis / novelty_matrix / experiment_plan / research_decision / research_log / manuscript_outline - 新增 evidence、results 数据与审计产物 - 新增 pyproject.toml、requirements、.gitignore
41 lines
2.2 KiB
Python
41 lines
2.2 KiB
Python
"""Single-device, batch-one reference training; no dataset auto-download."""
|
|
import argparse,json,random,os
|
|
from pathlib import Path
|
|
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8')
|
|
import numpy as np
|
|
import torch
|
|
from irstd.models import TinyUNet,segmentation_loss
|
|
from irstd.data import read_manifest,read_pair,sha
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser(); p.add_argument('--manifest',required=True);p.add_argument('--out',required=True)
|
|
p.add_argument('--epochs',type=int,default=100);p.add_argument('--seed',type=int,default=0)
|
|
p.add_argument('--loss',choices=['bce','dice','bce_dice','focal','iou'],default='bce_dice')
|
|
p.add_argument('--device',default='cuda'); a=p.parse_args()
|
|
if a.epochs<1: p.error('epochs must be positive')
|
|
random.seed(a.seed);np.random.seed(a.seed);torch.manual_seed(a.seed)
|
|
torch.use_deterministic_algorithms(True);torch.backends.cudnn.benchmark=False
|
|
rows=read_manifest(a.manifest);tr=[r for r in rows if r['split']=='train'];dv=[r for r in rows if r['split']=='dev']
|
|
if not tr or not dv: raise ValueError('Train and development splits required')
|
|
out=Path(a.out);out.mkdir(parents=True,exist_ok=False)
|
|
model=TinyUNet().to(a.device);opt=torch.optim.AdamW(model.parameters(),lr=1e-3,weight_decay=1e-4)
|
|
config=vars(a)|dict(manifest_sha256=sha(a.manifest),torch=str(torch.__version__),numpy=np.__version__,cuda=torch.version.cuda)
|
|
(out/'config.json').write_text(json.dumps(config,indent=2));best=float('inf')
|
|
def batch(r):
|
|
x,y=read_pair(r)
|
|
return torch.from_numpy(x)[None,None].to(a.device),torch.from_numpy(y.astype(np.float32))[None,None].to(a.device)
|
|
with (out/'train.jsonl').open('w') as f:
|
|
for epoch in range(a.epochs):
|
|
model.train();random.shuffle(tr);total=0.
|
|
for r in tr:
|
|
x,y=batch(r);opt.zero_grad();loss=segmentation_loss(model(x),y,a.loss)
|
|
loss.backward();opt.step();total+=loss.item()
|
|
model.eval();dev=0.
|
|
with torch.no_grad():
|
|
for r in dv:
|
|
x,y=batch(r);dev+=segmentation_loss(model(x),y,a.loss).item()
|
|
dev/=len(dv);f.write(json.dumps(dict(epoch=epoch,train_loss=total/len(tr),dev_loss=dev))+'\n');f.flush()
|
|
if dev<best:
|
|
best=dev;torch.save(dict(state_dict=model.state_dict(),config=config,epoch=epoch),out/'best.pt')
|
|
if __name__=='__main__':main()
|