103 lines
5.1 KiB
Python
103 lines
5.1 KiB
Python
"""Resumable KNP detector collection; fixed model/data/build signature, bounded processes."""
|
|
import argparse
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import fcntl
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import time
|
|
|
|
|
|
def digest(path):
|
|
h = hashlib.sha256()
|
|
with path.open('rb') as f:
|
|
for block in iter(lambda: f.read(1024 * 1024), b''):
|
|
h.update(block)
|
|
return h.hexdigest()
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('model', type=Path)
|
|
ap.add_argument('out', type=Path)
|
|
ap.add_argument('chunk', nargs='?', type=int, default=1000)
|
|
ap.add_argument('--data', type=Path, default=Path.home() / 'run/cifar/vehicles')
|
|
ap.add_argument('--build', type=Path, default=Path(__file__).resolve().parents[1] / 'build')
|
|
ap.add_argument('--backend', choices=['cpu','cuda'], default='cpu')
|
|
ap.add_argument('--jobs', type=int)
|
|
ap.add_argument('--limit-per-split', type=int, help='Diagnostic subset only; use a separate output directory')
|
|
a = ap.parse_args()
|
|
if a.jobs is None: a.jobs = int(os.environ.get("JOBS", "2" if a.backend == "cuda" else "12"))
|
|
if a.chunk < 1 or not 1 <= a.jobs <= 64 or (a.limit_per_split is not None and a.limit_per_split < 1):
|
|
ap.error('Invalid chunk/jobs/limit')
|
|
a.model, a.out, a.data, a.build = (p.resolve() for p in (a.model, a.out, a.data, a.build))
|
|
a.out.mkdir(parents=True, exist_ok=True)
|
|
lock = (a.out / '.lock').open('w')
|
|
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
executable = a.build / 'bin/altai_infer'
|
|
backend = a.build / 'bin' / ('libknp-altai-cuda-backend.so' if a.backend == 'cuda' else 'libknp-cpu-single-threaded-backend.so')
|
|
k, stride, threshold, classes, _ = map(int, (a.model / 'meta.txt').read_text().split())
|
|
jobs, sizes = [], {}
|
|
for split in ('train', 'train_flip', 'val'):
|
|
image_file = a.data / f'{split}_images.bin'
|
|
size = image_file.stat().st_size
|
|
if size % 3072:
|
|
raise ValueError(f'Invalid image file: {image_file}')
|
|
n = size // 3072
|
|
if a.limit_per_split:
|
|
n = min(n, a.limit_per_split)
|
|
sizes[split] = n
|
|
for first in range(0, n, a.chunk):
|
|
jobs.append((split, first, min(a.chunk, n - first)))
|
|
signature = {'detector': [k, stride, threshold, classes], 'data_sizes': sizes, 'chunk': a.chunk,
|
|
'files': {str(p): digest(p) for p in [executable, backend, a.model/'det_w.bin', a.model/'det_b.bin',
|
|
*(a.data/f'{s}_images.bin' for s in sizes)]}}
|
|
settings = a.out / 'collection-settings.json'
|
|
if settings.exists() and json.loads(settings.read_text()) != signature:
|
|
raise ValueError('Different model/data/build/chunk: use a new output directory')
|
|
settings.write_text(json.dumps(signature, indent=2) + '\n')
|
|
def valid(job):
|
|
split, first, count = job
|
|
prefix = a.out / f'{split}_{first}'
|
|
output = Path(str(prefix) + '_counts.bin')
|
|
return Path(str(prefix) + '.done').exists() and output.exists() and output.stat().st_size == count * 4 * k * 4
|
|
finished = sum(n for s, first, n in jobs if valid((s, first, n)))
|
|
total = sum(sizes.values())
|
|
started = time.time()
|
|
def progress(failed=False, error=None):
|
|
state = {'current': finished, 'total': total, 'done': finished == total and not failed,
|
|
'failed': failed, 'error': error, 'started_unix': started, 'updated_unix': time.time()}
|
|
temporary = a.out/'progress.tmp'
|
|
temporary.write_text(json.dumps(state) + '\n')
|
|
temporary.replace(a.out/'progress.json')
|
|
def one(job):
|
|
split, first, count = job
|
|
prefix = a.out / f'{split}_{first}'
|
|
env = dict(os.environ, LD_LIBRARY_PATH=str(a.build/'bin') + ':' + os.environ.get('LD_LIBRARY_PATH',''))
|
|
with Path(str(prefix) + '.log').open('w') as log:
|
|
result = subprocess.run([str(executable), str(a.data/f'{split}_images.bin'), str(first), str(count),
|
|
str(a.model), str(backend), str(prefix)], stdout=log, stderr=subprocess.STDOUT, env=env)
|
|
output = Path(str(prefix) + '_counts.bin')
|
|
if result.returncode or not output.exists() or output.stat().st_size != count * 4 * k * 4:
|
|
raise RuntimeError(f'{split}_{first}: failed or truncated output; inspect its log')
|
|
Path(str(prefix) + '.done').write_text(json.dumps({'count':count, 'completed_unix':time.time()}) + '\n')
|
|
return count
|
|
progress()
|
|
errors = []
|
|
with ThreadPoolExecutor(max_workers=a.jobs) as pool:
|
|
futures = {pool.submit(one, job): job for job in jobs if not valid(job)}
|
|
for future in as_completed(futures):
|
|
try:
|
|
finished += future.result()
|
|
except Exception as error:
|
|
errors.append(str(error))
|
|
progress(bool(errors), '; '.join(errors) if errors else None)
|
|
print(f'{finished}/{total}' + (' FAILED '+errors[-1] if errors else ''), flush=True)
|
|
if errors:
|
|
raise RuntimeError('; '.join(errors))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|