File size: 3,339 Bytes
6a5bb7e | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Model Soup: average weights of multiple seed models → single model.
Zero inference cost vs single model, near-ensemble performance.
"""
import sys, os
import copy
import numpy as np
PROJECT_DIR = '/home/user/goat'
os.chdir(PROJECT_DIR)
sys.path.insert(0, PROJECT_DIR)
import torch
from ultralytics import YOLO
from tqdm import tqdm
def soup_models(model_paths, output_path):
"""Average weights of multiple YOLO models."""
models = [YOLO(p) for p in model_paths]
base = models[0]
# Collect state dicts
state_dicts = []
for m in models:
sd = m.model.state_dict()
state_dicts.append(sd)
# Average
avg_sd = {}
for key in state_dicts[0].keys():
if any(kw in key.lower() for kw in ['num_batches_tracked', 'running_mean', 'running_var']):
# Use first model's BN stats
avg_sd[key] = state_dicts[0][key].clone()
elif 'bn' in key.lower() or 'batch_norm' in key.lower():
# Average BN stats
avg_sd[key] = sum(sd[key] for sd in state_dicts) / len(state_dicts)
else:
avg_sd[key] = sum(sd[key] for sd in state_dicts) / len(state_dicts)
base.model.load_state_dict(avg_sd)
base.model.eval()
base.save(output_path)
print(f'Model Soup saved to {output_path}')
return output_path
def evaluate_model(model_path, name):
"""Quick eval on val set."""
model = YOLO(model_path)
results = model.val(data='Data/Detection_dataset/dataset.yaml', imgsz=1536, batch=2, verbose=False)
return float(results.box.map50), float(results.box.map)
def main():
model_paths = [
'runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt',
'runs/detect/Detection_experiments/v12_seed_42/weights/best.pt',
'runs/detect/Detection_experiments/v12_seed_123/weights/best.pt',
]
# Verify all exist
valid = [p for p in model_paths if os.path.exists(p)]
print(f'Models available: {len(valid)}/{len(model_paths)}')
for p in valid:
print(f' {p}')
# Eval individual models
print(f'\n{"="*50}')
print('Individual models:')
singles = {}
for p in valid:
name = p.split('/')[-3]
m50, m5095 = evaluate_model(p, name)
singles[name] = (m50, m5095)
print(f' {name}: mAP50={m50:.4f}, mAP50-95={m5095:.4f}')
# Model Soup
print(f'\n{"="*50}')
print('Model Soup experiments:')
# Try all combinations
import itertools
for r in range(2, len(valid) + 1):
for combo in itertools.combinations(range(len(valid)), r):
paths = [valid[i] for i in combo]
names = [p.split('/')[-3] for p in paths]
label = ' + '.join(n.split('_')[0] + '_' + n.split('_')[-1][:4] for n in names)
output = f'runs/soup/soup_{len(paths)}models.pt'
os.makedirs('runs/soup', exist_ok=True)
soup_models(paths, output)
m50, m5095 = evaluate_model(output, label)
delta = m5095 - singles[names[0]][1]
print(f' {label}: mAP50={m50:.4f}, mAP50-95={m5095:.4f} (d={delta:+.4f} vs {names[0]})')
# WBF baseline for comparison
print(f'\n{"="*50}')
print('WBF Ensemble (for comparison):')
print(' (Run eval_ensemble.py for full WBF results)')
if __name__ == '__main__':
main()
|