Digital Drosophila Archive
62 lines
2.9 KiB
Python
62 lines
2.9 KiB
Python
"""Rerun old deterministic-v2 references without changing their semantics."""
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import numpy as np
|
|
import pandas as pd
|
|
import yaml
|
|
from context import ROOT, EXP, check, save, bit_equal
|
|
|
|
|
|
def preflight():
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = 'GPU-79f3a85d-adce-7833-7003-6488140ca596'
|
|
os.environ['CUDA_HOME'] = '/usr/local/cuda'
|
|
os.environ['PATH'] = str(ROOT / 'models/male-cns-lif/.venv/bin') + ':/usr/local/cuda/bin:' + os.environ['PATH']
|
|
os.environ['MAX_JOBS'] = '2'
|
|
os.environ['TORCH_CUDA_ARCH_LIST'] = '8.9'
|
|
for port in (8188, 8189, 8190):
|
|
with urllib.request.urlopen(f'http://192.168.10.51:{port}/queue', timeout=5) as response:
|
|
queue = json.load(response)
|
|
assert not queue['queue_running'] and not queue['queue_pending'], 'Comfy busy: defer, do not evict'
|
|
|
|
|
|
def main():
|
|
check(); preflight()
|
|
assert not (EXP / 'regression/frozen-five.json').exists()
|
|
runtime_path = ROOT / 'models/male-cns-lif/deterministic-runtime-v2'
|
|
config = yaml.safe_load((runtime_path / 'phase4p.yaml').read_text())
|
|
sys.path.insert(0, str(ROOT / 'models/male-cns-lif'))
|
|
sys.path.insert(0, str(runtime_path))
|
|
import torch
|
|
import phase4p
|
|
assert torch.cuda.mem_get_info()[0] > 2 * 1024**3
|
|
torch.cuda.reset_peak_memory_stats()
|
|
started = time.perf_counter()
|
|
runtime = phase4p.CanonicalRuntime(config, runtime_path, EXP)
|
|
rows = []
|
|
for spec in config['reference_trials']:
|
|
preflight()
|
|
reference_dir = ROOT / 'experiments/phase4p-deterministic-runtime/phase4c-v2/reference' / spec['name']
|
|
events = pd.read_parquet(reference_dir / 'spike-events.parquet')
|
|
reference = dict(step=events.step.to_numpy(), neuron_idx=events.neuron_idx.to_numpy(),
|
|
counts=np.load(reference_dir / 'per-neuron-spike-counts.int32.npy'),
|
|
v=np.load(reference_dir / 'final-voltage.float32.npy'), g=np.load(reference_dir / 'final-g.float32.npy'))
|
|
result = runtime.run(np.flatnonzero(runtime.pops[spec['population']]), spec['rate_hz'], spec['seed'], spec['duration_ms'])
|
|
row = dict(trial=spec['name'], exact={k: bit_equal(v, result[k]) for k, v in reference.items()},
|
|
dtype_shape={k: dict(reference_dtype=str(v.dtype), actual_dtype=str(result[k].dtype),
|
|
reference_shape=list(v.shape), actual_shape=list(result[k].shape)) for k, v in reference.items()},
|
|
wall_seconds=result['metrics']['wall_seconds'])
|
|
rows.append(row)
|
|
print(row, flush=True)
|
|
passed = len(rows) == 5 and all(all(row['exact'].values()) for row in rows)
|
|
save(EXP / 'regression/frozen-five.json', dict(pass_all=passed, trials=rows,
|
|
wall_seconds=time.perf_counter() - started, peak_allocated_vram=int(torch.cuda.max_memory_allocated()),
|
|
scope='Old frozen deterministic-v2 reference only, not D18 equality to old dynamics'))
|
|
assert passed
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|