| |
| """Solve or train, save, reload, and sample the Lecture 7 methods.""" |
| import argparse |
| import inspect |
| import json |
| import platform |
| import time |
| from pathlib import Path |
| import torch |
| import numpy as np |
| from ot_sbm_examples import (serial, transport_example, finite_bridge_example, |
| discrete_imf_example, ctmc_bridge_example, gaussian_example, |
| reward_bridge_example, branching_example, entangled_geometry_example) |
| from learned_bridges import METHODS as LEARNED, save_artifact |
| from discrete_learning import fit_ddsbm, fit_csbm |
| from sampling import generate |
|
|
| ROOT = Path(__file__).resolve().parent |
| EXACT = {'ot': transport_example, 'sinkhorn': transport_example, |
| 'finite-sb': finite_bridge_example, 'discrete-imf': discrete_imf_example, |
| 'ctmc-sb': ctmc_bridge_example, 'gaussian-sb': gaussian_example, |
| 'reward-tilt': reward_bridge_example, 'branch-mass': branching_example, |
| 'cone-geometry': entangled_geometry_example} |
| LEARNED = {**LEARNED, 'ddsbm': fit_ddsbm, 'csbm': fit_csbm} |
| METHODS = list(EXACT) + list(LEARNED) |
| SEEDS = dict(dsb=7, dsbm=8, sf2m=9, tr2d2=10, branch=11, entangled=12, ddsbm=13, csbm=14) |
|
|
| def write_json(path, value): |
| Path(path).write_text(json.dumps(serial(value), indent=2, allow_nan=False)+'\n') |
|
|
| def parser(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--method', choices=METHODS) |
| p.add_argument('--mode', choices=['train-sample', 'train', 'sample'], default='train-sample') |
| p.add_argument('--out', help='Run directory, default lecture_7/outputs/METHOD') |
| p.add_argument('--seed', type=int, help='Training seed; sampling uses this seed plus 100') |
| p.add_argument('--samples', type=int, default=256) |
| p.add_argument('--sample-steps', type=int, help='Override the inference grid when the sampler supports it') |
| p.add_argument('--train-steps', type=int, help='Gradient updates for sf2m, branch, ddsbm, or csbm') |
| p.add_argument('--rounds', type=int, help='Forward/reverse fitting cycles for dsb or dsbm') |
| p.add_argument('--epochs', type=int, help='Replay/CE epochs for tr2d2 or entangled') |
| p.add_argument('--batch-size', type=int, help='Training batch size for methods with a batch argument') |
| p.add_argument('--searches', type=int, help='MCTS rollouts per TR2-D2 epoch') |
| p.add_argument('--quick', action='store_true', help='Short execution check; not a quality experiment') |
| return p |
|
|
| def main(argv=None): |
| p = parser() |
| args = p.parse_args(argv) |
| for name in ['samples','sample_steps','train_steps','rounds','epochs','batch_size','searches']: |
| value = getattr(args, name) |
| if value is not None and value < 1: |
| p.error(name.replace('_','-')+' must be positive') |
| if args.mode == 'sample' and args.out is None: |
| p.error('Sample mode requires --out for its checkpoint.') |
| if args.mode != 'sample' and args.method is None: |
| args.method = 'sinkhorn' |
| out = Path(args.out) if args.out else ROOT/'outputs'/args.method |
| out.mkdir(parents=True, exist_ok=True) |
| began = time.perf_counter() |
| if args.mode == 'sample': |
| state = torch.load(out/'checkpoint.pt', map_location='cpu', weights_only=True) |
| if args.method is not None and args.method != state['method']: |
| p.error('The requested method differs from the saved checkpoint.') |
| seed = args.seed if args.seed is not None else json.loads((out/'config.json').read_text())['seed'] |
| else: |
| seed = args.seed if args.seed is not None else SEEDS.get(args.method, 6270) |
| config = {'method': args.method, 'seed': seed, 'mode': args.mode, |
| 'quick': args.quick, 'device': 'cpu', 'python': platform.python_version(), |
| 'torch': torch.__version__, 'numpy': np.__version__} |
| if args.method in EXACT: |
| if any(x is not None for x in [args.train_steps,args.rounds,args.epochs,args.batch_size,args.searches]): |
| p.error('This method is a numerical solver; it has no neural training parameters.') |
| report = EXACT[args.method]() |
| steps = 2 if args.method == 'finite-sb' else 100 |
| save_artifact(out/'checkpoint.pt', method=args.method, result=serial(report), steps=steps) |
| config['training'] = 'Exact or explicitly enumerated numerical calculation' |
| else: |
| fn = LEARNED[args.method] |
| signature = inspect.signature(fn).parameters |
| settings = {'seed': seed, 'checkpoint': out/'checkpoint.pt'} |
| if args.quick: |
| for name, value in dict(updates=20, rounds=2, epochs=1, batch=256, searches=40).items(): |
| if name in signature: |
| settings[name] = value |
| for arg, param in [('train_steps','updates'),('rounds','rounds'),('epochs','epochs'),('batch_size','batch'),('searches','searches')]: |
| value = getattr(args, arg) |
| if value is not None: |
| if param not in signature: |
| p.error('--'+arg.replace('_','-')+' does not apply to '+args.method) |
| settings[param] = value |
| config['training'] = {k: settings.get(k, v.default) for k,v in signature.items() if k != 'checkpoint'} |
| report = fn(**settings) |
| write_json(out/'config.json', config) |
| write_json(out/'report.json', report) |
| write_json(out/'losses.json', report.get('loss_samples', report.get('history', []))) |
| state = torch.load(out/'checkpoint.pt', map_location='cpu', weights_only=True) |
| if args.mode == 'train': |
| print(json.dumps({'method': args.method, 'mode': args.mode, 'out': str(out), 'seconds': time.perf_counter()-began})) |
| return report |
| output = generate(state, args.samples, seed+100, args.sample_steps) |
| stem = 'resampled' if args.mode == 'sample' else 'samples' |
| write_json(out/(stem+'.json'), output) |
| (out/(stem+'.txt')).write_text('\n'.join(json.dumps(row) for row in output['values'])+'\n') |
| write_json(out/('sample_report.json' if args.mode == 'sample' else 'generation_report.json'), output['report']) |
| print(json.dumps({'method':state['method'],'mode':args.mode,'out':str(out),'seconds':round(time.perf_counter()-began,3)}),flush=True) |
| return output |
|
|
| if __name__ == '__main__': |
| main() |
|
|
