File size: 3,533 Bytes
cd9b2d8 | 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 | #!/usr/bin/env python3
"""Evaluate saved best checkpoints on their existing selection-validation split."""
import os
import torch
from study_paths import ROOT, RELEASE, GEMINI, CODE, read, write
from cache_dataset import FeatureData
from probe_model import Probe
from train_probes import evaluate
from evaluate_whisper import LayeredMultiTaskWhisper, GeminiDataset, HoldoutLoader, forward
def run_probe(model_id, variant, phase, device, classes):
out = ROOT / 'probes' / phase / model_id / variant
target = out / 'validation_metrics.json'
if target.exists():
return
saved = torch.load(out / 'best.pt', map_location='cpu', weights_only=False)
domains = ['gemini'] if phase == 'gemini' else ['ladder', 'p3_train', 'p3_validation', 'p3_test']
data = FeatureData(model_id, domains, 'validation', phase)
data.projection = saved['pca']
model = Probe(len(classes), variant == 'linear').to(device)
model.load_state_dict(saved['model'], strict=True)
del saved
metrics = evaluate(model, data, device, torch.ones(len(classes), device=device), True)
metrics.update(split='validation', phase=phase, domain='Flash' if phase == 'gemini' else 'Ladder/P3',
split_scope='Existing checkpoint-selection validation split; not a fresh independent test')
write(target, metrics)
print('VALIDATED', phase, model_id, variant, 'clips', metrics['clips'], flush=True)
del model, data
torch.cuda.empty_cache()
def run_whisper(size, phase, device, classes):
out = ROOT / 'whisper' / phase / size
target = out / 'validation_metrics.json'
if target.exists():
return
cfg = read(CODE / 'gemini_finetune/configs' / ('whisper_' + size + '.json'))
checkpoint = cfg['checkpoint'] if phase == 'legacy' else GEMINI / 'training' / ('whisper_' + size) / 'best.pt'
norm = read(RELEASE / 'training_normalization.json')
model = LayeredMultiTaskWhisper(RELEASE / size, len(classes), initialize_pretrained=False).to(device)
model.load_state_dict(torch.load(checkpoint, map_location='cpu', weights_only=False)['model'], strict=True)
data = GeminiDataset(GEMINI / 'prepared', 'validation')
metrics = evaluate(model, data, device, torch.ones(len(classes), device=device), True,
HoldoutLoader(RELEASE / size, norm), forward)
metrics.update(split='validation', phase=phase, domain='Flash',
split_scope='Same Flash validation audio before/after; used to select tuned checkpoint')
write(target, metrics)
print('VALIDATED Whisper', size, phase, 'clips', metrics['clips'], flush=True)
del model, data
torch.cuda.empty_cache()
def main():
rank = int(os.environ.get('LOCAL_RANK', 0)); device = torch.device('cuda', rank)
torch.cuda.set_device(device); torch.set_num_threads(2)
classes = read(RELEASE / 'classes.json')['names']
size, phase = [('base', 'legacy'), ('base', 'gemini'), ('small', 'legacy'), ('small', 'gemini')][rank]
run_whisper(size, phase, device, classes)
jobs = [(m['id'], v, p) for p in ('legacy', 'gemini')
for m in read(ROOT / 'study.json')['models'] for v in ('linear', 'mlp')]
for index, job in enumerate(jobs):
if index % 4 == rank:
run_probe(*job, device, classes)
write(ROOT / 'consolidated_benchmark' / ('VALIDATION_COMPLETE-rank' + str(rank) + '.json'),
{'rank': rank, 'probe_configurations': len(jobs[rank::4]), 'whisper_configurations': 1})
if __name__ == '__main__':
main()
|