File size: 4,400 Bytes
954bb80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import datetime
import json
import subprocess
import sys
import time
from pathlib import Path
ROOT=Path('/home/ubuntu/tsr')
OUT=ROOT/'research_round'
PLAN=[{'method':'l2sp','epochs':100,'lr':1e-4,'init':'pretrained','patience':30},
{'method':'sam','epochs':20,'lr':1e-5,'init':'validation_winner','patience':0},
{'method':'contrast','epochs':100,'lr':1e-4,'init':'pretrained','patience':30},
{'method':'consistency','epochs':100,'lr':1e-4,'init':'pretrained','patience':30},
{'method':'swa','epochs':20,'lr':1e-5,'init':'validation_winner','patience':0}]
(OUT/'plan.json').write_text(json.dumps(PLAN,indent=2))
def command(trial,smoke=False):
method=trial['method']
cmd=[sys.executable,'-u',str(ROOT/'finetune_research.py'),'--pretrain_pth',str(ROOT/'latest.utae.tar'),
'--research',method,'--tag',f'32X32_research_{method}', '--out_dir',str(OUT/('smoke' if smoke else 'checkpoints')/method),
'--epochs',str(trial['epochs']),'--base_lr',str(trial['lr']),'--encoder_lr_mult','0.5',
'--batch','2','--crop_size','32','--fixed_crops','--mask_ce_invalid','--resume','none',
'--patience',str(trial['patience']),'--warmup_epochs', '0' if trial['init']=='validation_winner' else '3']
if trial['init']=='validation_winner':
selection=OUT/'initialization.json'
initial=json.loads(selection.read_text())['checkpoint'] if selection.exists() else str(ROOT/'checkpoints/finetune_v2/h3_mask_ce_invalid/model_best_raw.tar')
cmd+=['--init_finetuned',initial]
if not smoke and method == 'contrast':
cmd+=['--resume','auto']
if smoke:
# Exercise consistency at epoch 3 and SWA averaging/BN calibration at epoch 5.
cmd+=['--epochs',str(6 if method=='swa' else 4 if method=='consistency' else 1),'--limit_batches','1','--patience','0']
return cmd
def run(trial,smoke=False):
name=trial['method']; key=('smoke_' if smoke else '')+name
statusfile=OUT/'queue_status.json'
status=json.loads(statusfile.read_text()) if statusfile.exists() else {}
status[key]={'status':'running','started':datetime.datetime.now(datetime.timezone.utc).isoformat(),'command':command(trial,smoke)}
statusfile.write_text(json.dumps(status,indent=2))
start=time.monotonic()
with (OUT/f'{key}.log').open('w') as stream:
result=subprocess.run(command(trial,smoke),cwd=ROOT,stdout=stream,stderr=subprocess.STDOUT)
status[key].update(status='completed' if result.returncode==0 else 'failed',seconds=time.monotonic()-start,returncode=result.returncode)
statusfile.write_text(json.dumps(status,indent=2))
if result.returncode: raise RuntimeError(f'{key} failed; inspect {key}.log')
print(key,'completed',status[key]['seconds'],'seconds',flush=True)
def main():
if '--smoke' in sys.argv:
for trial in PLAN: run(trial,True)
return
import hashlib
manifest=json.loads((OUT/'protocol_manifest.json').read_text())
for name,digest in manifest.items():
if hashlib.sha256((ROOT/name).read_bytes()).hexdigest()!=digest:
raise RuntimeError(f'Protocol or pretrained file changed: {name}')
print('Waiting for existing improvement round to finish; no competing GPU training.',flush=True)
while 'Improvement round complete' not in (ROOT/'improvement_round/service.log').read_text():
state=subprocess.run(['systemctl','--user','show','tsr-improvement-round.service','--property=ActiveState','--value'],capture_output=True,text=True)
if state.returncode==0 and state.stdout.strip() in ['failed','inactive']:
raise RuntimeError('Prior improvement round stopped before completion')
time.sleep(30)
(OUT/'initialization.json').write_text((ROOT/'improvement_round/selected_model.json').read_text())
status=json.loads((OUT/'queue_status.json').read_text())
for trial in PLAN:
if status.get('smoke_'+trial['method'],{}).get('status')!='completed': run(trial,True)
for trial in PLAN:
current=json.loads((OUT/'queue_status.json').read_text())
if current.get(trial['method'],{}).get('status') == 'completed':
print('Retaining completed trial:',trial['method'],flush=True)
continue
run(trial)
print('Research queue completed. Test evaluations run in the separate test service.',flush=True)
if __name__=='__main__': main()
|