| """Evaluate CIFAR-10 checkpoints with independent per-GPU seeds.""" |
| import argparse |
| import csv |
| import hashlib |
| import io |
| import json |
| import os |
| from pathlib import Path |
| import subprocess |
| import sys |
| import time |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
| |
| CONFIGS = { |
| "baseline": "baseline_ar.yaml", |
| "tail_dropout": "tail_dropout_ar.yaml", |
| "flow_matching": "flow_matching_ar.yaml", |
| "flow_matching_tail_dropout": "flow_matching_tail_dropout_ar.yaml", |
| "flow_matching_tied_tail_dropout": "flow_matching_tied_tail_dropout_ar.yaml", |
| "war": "war_ar.yaml" |
| } |
|
|
|
|
| def sha256(path): |
| digest = hashlib.sha256() |
| with open(path, 'rb') as f: |
| for block in iter(lambda: f.read(8 * 1024 * 1024), b''): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def check_free(gpus): |
| output = subprocess.check_output( |
| ['nvidia-smi', '--query-gpu=index,memory.used,utilization.gpu', |
| '--format=csv,noheader,nounits'], text=True) |
| devices = {int(row[0]): (int(row[1]), int(row[2])) |
| for row in csv.reader(io.StringIO(output))} |
| busy = {gpu: devices.get(gpu) for gpu in gpus |
| if gpu not in devices or devices[gpu][0] > 100 or devices[gpu][1] > 0} |
| if busy: |
| raise RuntimeError(f'Requested GPUs are not free: {busy}') |
|
|
|
|
| def save_summary(output, rows): |
| (output / 'results.json').write_text(json.dumps(rows, indent=2) + '\n') |
| if rows: |
| with (output / 'results.csv').open('w', newline='') as f: |
| writer = csv.DictWriter(f, fieldnames=list(rows[0])) |
| writer.writeheader() |
| writer.writerows(rows) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--gpus', default='0,1,2,3,4,5,6,7') |
| parser.add_argument('--output', default='outputs/cifar10-nonlegacy') |
| parser.add_argument('--methods', nargs='+', choices=CONFIGS, |
| help='Optional subset, e.g. baseline war') |
| parser.add_argument('--resume', action='store_true', help='Reuse matching completed evaluations') |
| args = parser.parse_args() |
| gpus = [int(x) for x in args.gpus.split(',')] |
| if not gpus or len(gpus) != len(set(gpus)): |
| parser.error('GPU IDs must be nonempty and distinct') |
| configs = {key: ROOT / 'configs/cifar10' / name for key, name in CONFIGS.items()} |
| for config in configs.values(): |
| if not config.is_file(): |
| raise FileNotFoundError(config) |
| if args.methods: |
| if set(args.methods) - configs.keys(): |
| parser.error('Unknown method') |
| configs = {key: path for key, path in configs.items() if key in args.methods} |
| checkpoints = {key: ROOT / 'checkpoints' / f'{key}.pt' for key in configs} |
| for checkpoint in checkpoints.values(): |
| if not checkpoint.is_file(): |
| raise FileNotFoundError(f'Place the downloaded checkpoint at {checkpoint}') |
| reference = ROOT / 'refs/cifar10-32x32_local.npz' |
| reference_hash = sha256(reference) |
| if reference_hash != '5493968c4f29b2878f64d6a640351f377e32e86553db56eca8147a0ae50d93f2': |
| raise ValueError('The CIFAR-10 FID reference checksum does not match') |
| output = Path(args.output).resolve() |
| output.mkdir(parents=True, exist_ok=True) |
| env = os.environ.copy() |
| env.update(CUDA_VISIBLE_DEVICES=','.join(map(str, gpus)), PYTHONDONTWRITEBYTECODE='1', |
| WANDB_MODE='disabled', OMP_NUM_THREADS='8', MKL_NUM_THREADS='8', TF_CPP_MIN_LOG_LEVEL='2') |
| rows = [] |
| for key, config in configs.items(): |
| checkpoint = checkpoints[key] |
| checkpoint_hash, config_hash = sha256(checkpoint), sha256(config) |
| case = output / key |
| result_path = case / 'metrics.json' |
| if result_path.exists() and not args.resume: |
| raise FileExistsError(result_path) |
| if not result_path.exists(): |
| check_free(gpus) |
| case.mkdir(exist_ok=True) |
| cmd = [sys.executable, '-m', 'torch.distributed.run', '--standalone', |
| f'--nproc-per-node={len(gpus)}', 'evaluate.py', '--config', str(config), |
| '--checkpoint', str(checkpoint), '--num-samples', '50000', '--batch-size', '64', |
| '--seed', '42', '--no-legacy', '--output', str(case), '--fid-reference', str(reference)] |
| request = dict(command=cmd, checkpoint_sha256=checkpoint_hash, |
| config_sha256=config_hash, legacy=False) |
| (case / 'request.json').write_text(json.dumps(request, indent=2) + '\n') |
| print(f'START {key}', flush=True) |
| started = time.monotonic() |
| with (case / 'run.log').open('w') as log: |
| subprocess.run(cmd, cwd=ROOT, env=env, stdout=log, stderr=subprocess.STDOUT, check=True) |
| print(f'DONE {key} in {time.monotonic() - started:.1f}s', flush=True) |
| result = json.loads(result_path.read_text()) |
| request = json.loads((case / 'request.json').read_text()) |
| if request['checkpoint_sha256'] != checkpoint_hash: |
| raise ValueError(f'{key}: checkpoint changed since the recorded evaluation') |
| expected = dict(legacy=False, num_samples=50000, seed=42, world_size=len(gpus), |
| batch_size=64, ar_ema=True, precision='32-true', decoder='single_step', |
| config_sha256=config_hash, reference_sha256=reference_hash) |
| if any(result.get(name) != value for name, value in expected.items()): |
| raise ValueError(f'{key}: recorded evaluation settings do not match') |
| row = dict(method=key, nonlegacy_gfid=result['fid'], unique_images=result['unique_images'], |
| num_samples=result['num_samples'], legacy=False, |
| checkpoint_sha256=checkpoint_hash, config_sha256=config_hash, |
| reference_sha256=reference_hash) |
| rows.append(row) |
| save_summary(output, rows) |
| print(f'RESULT {key}: {result["fid"]:.8f}', flush=True) |
| print(f'COMPLETE: {len(rows)} non-legacy evaluations', flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
