Blame for train.py

#!/usr/bin/env python3
"""Train ALL 15 operation families from independent random initialization.
No pretrained checkpoint is read. Validation controls early stopping and blending;
test/OOD metrics are written only after selection. The five training stages run sequentially to limit peak GPU memory use.
"""
import argparse, json, os, subprocess, sys, time
from pathlib import Path
ROOT=Path(__file__).resolve().parent
sys.path.insert(0,str(ROOT/'src'))
from config import VARIANT
from prepare_data import prepare

def main():
 prepare()
 p=argparse.ArgumentParser()
 p.add_argument('--device',choices=('auto','cpu','cuda'),default='auto')
 p.add_argument('--quick',action='store_true',help='Pipeline check ONLY: undertrained; NOT publishable metrics')
 p.add_argument('--stage',choices=('all','manufacturing','robotics','gmaw','rsw','tig','report'),default='all')
 p.add_argument('--seed',type=int,default=20261018)
 p.add_argument('--batch',type=int,default=8,help='manufacturing Time Transformer batch; start with 4 or 8 on limited GPU memory')
 args=p.parse_args()
 print('Manufacturing-MoE',VARIANT,'device',args.device,'quick',args.quick,flush=True)
 env=os.environ.copy();env.update(OMP_NUM_THREADS='1',OPENBLAS_NUM_THREADS='1',MKL_NUM_THREADS='1',PYTHONUNBUFFERED='1')
 epoch=2 if args.quick else None
 cmds={
  'manufacturing':['src/train.py','--output',str(ROOT/'reports'/'trained'),'--device',args.device,'--seed',str(args.seed),'--batch',str(args.batch),
    '--epochs-time',str(epoch or 65),'--epochs-energy',str(epoch or 120),'--epochs-machine',str(epoch or 300),
    '--patience',str(2 if args.quick else 15),'--patience-energy',str(2 if args.quick else 25), '--smoke-max-per-group',str(8 if args.quick else 0)],
  'robotics':['src/train_robotics.py','--device',args.device,'--seed',str(args.seed+1),'--epochs',str(epoch or 45),'--patience',str(2 if args.quick else 12)],
  'gmaw':['src/train_arc_welding.py','--device',args.device,'--seed',str(args.seed+2),'--epochs',str(epoch or 260),'--patience',str(2 if args.quick else 55)],
  'rsw':['src/train_welding.py','--process','rsw','--device',args.device,'--seed',str(args.seed+3),'--epochs',str(epoch or 260),'--patience',str(2 if args.quick else 55)],
  'tig':['src/train_welding.py','--process','tig','--device',args.device,'--seed',str(args.seed+4),'--epochs',str(epoch or 260),'--patience',str(2 if args.quick else 55)],
 }
 chosen=list(cmds) if args.stage=='all' else ([args.stage] if args.stage!='report' else [])
 completed=[]
 for stage in chosen:
  start=time.time();print('START',stage,flush=True)
  subprocess.run([sys.executable,*cmds[stage]],cwd=ROOT,env=env,check=True)
  completed.append({'stage':stage,'elapsed_seconds':time.time()-start});print('DONE',stage,flush=True)
 if chosen:(ROOT/'reports').mkdir(exist_ok=True);(ROOT/'reports'/'stage_durations.json').write_text(json.dumps(completed,indent=2))
 if args.stage in ('all','report'):
  subprocess.run([sys.executable,'src/make_report.py',*(['--quick'] if args.quick else [])],cwd=ROOT,check=True)
 print('Training workflow ended. Reports: reports/training_report.md and reports/metrics.csv')
if __name__=='__main__':main()