mirror of
https://github.com/jung-geun/PSO.git
synced 2026-09-20 14:11:48 +09:00
Migrate the package and examples to the tensor-native PyTorch implementation, add benchmark evidence, and add the guarded post-training convergence protocol with TensorBoard progress monitoring and hash-verified recovery. Constraint: Preserve one-shot official-test sealing and auditable research artifacts Rejected: Commit local .omc runs and downloaded datasets | multi-gigabyte runtime state is machine-local Confidence: high Scope-risk: broad Not-tested: Production CUDA run on pieroot-server
236 lines
8.8 KiB
Python
236 lines
8.8 KiB
Python
"""
|
|
Focused Offline Unit Tests for MNIST Deep PSO Methods (Protocol MNIST-PSO-RAW-V5 1.0.0)
|
|
|
|
Tests:
|
|
1. Split balance and nested subset inclusion invariant (I_2k subset of I_10k subset of I_50k).
|
|
2. Latent transform exactness (z_0 = 0 -> theta_0), antithetic symmetry (even swarm pairing), and subspace dimensions.
|
|
3. Stage transition pbest re-evaluation, gbest rebuild, and exact query/sample accounting.
|
|
4. Validation-only pilot and elite ensemble selection with greedy disagreement.
|
|
5. CLI argument validation.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Insert repo test/ path so deep_pso_methods can be imported
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "test"))
|
|
|
|
import math
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from deep_pso_methods import (
|
|
CompactCNN,
|
|
LatentTransform,
|
|
build_nested_stratified_subsets,
|
|
evaluate_latent_batch,
|
|
evaluate_probabilistic_metrics,
|
|
make_compact_cnn,
|
|
run_latent_pso,
|
|
select_diverse_candidates,
|
|
validate_cli_args,
|
|
build_parser,
|
|
)
|
|
|
|
|
|
def test_split_balance_and_nesting():
|
|
"""Verify stratified split balance and nested index inclusion: I_2k subset of I_10k subset of I_50k."""
|
|
N = 50000
|
|
num_classes = 10
|
|
samples_per_class = N // num_classes
|
|
y_search = torch.cat([torch.full((samples_per_class,), c, dtype=torch.long) for c in range(num_classes)])
|
|
|
|
split_seed = 20260902
|
|
nested_subsets = build_nested_stratified_subsets(
|
|
y_search=y_search,
|
|
subset_sizes=[2000, 10000, 50000],
|
|
subset_seed=split_seed,
|
|
)
|
|
|
|
idx_2k = nested_subsets[2000]
|
|
idx_10k = nested_subsets[10000]
|
|
idx_50k = nested_subsets[50000]
|
|
|
|
assert len(idx_2k) == 2000
|
|
assert len(idx_10k) == 10000
|
|
assert len(idx_50k) == 50000
|
|
|
|
set_2k = set(idx_2k.numpy().tolist())
|
|
set_10k = set(idx_10k.numpy().tolist())
|
|
set_50k = set(idx_50k.numpy().tolist())
|
|
|
|
# Strict nesting: I_2k subset of I_10k subset of I_50k
|
|
assert set_2k.issubset(set_10k), "I_2k must be a strict subset of I_10k"
|
|
assert set_10k.issubset(set_50k), "I_10k must be a strict subset of I_50k"
|
|
|
|
# Exact stratification (equal counts per class)
|
|
y_2k = y_search[idx_2k].numpy()
|
|
y_10k = y_search[idx_10k].numpy()
|
|
|
|
counts_2k = np.bincount(y_2k, minlength=10)
|
|
counts_10k = np.bincount(y_10k, minlength=10)
|
|
|
|
for c in range(num_classes):
|
|
assert counts_2k[c] == 200, f"Class {c} in 2k subset must have 200 samples; got {counts_2k[c]}"
|
|
assert counts_10k[c] == 1000, f"Class {c} in 10k subset must have 1000 samples; got {counts_10k[c]}"
|
|
|
|
|
|
def test_latent_transform_exactness_and_antithetic_symmetry():
|
|
"""Verify transform exactness (z_0 = 0 -> theta_0), even swarm pairing, zero centroid, and subspace shapes."""
|
|
device = torch.device("cpu")
|
|
base_model = make_compact_cnn(seed=41).to(device)
|
|
base_vec = torch.cat([p.detach().view(-1) for p in base_model.parameters()])
|
|
|
|
dims = [290, 1024, 4096, "full"]
|
|
swarm_size = 30 # Even swarm size
|
|
|
|
for d in dims:
|
|
transform = LatentTransform(base_model, latent_dim=d, device=device)
|
|
Z = transform.init_swarm(swarm_size=swarm_size, seed=91)
|
|
|
|
# 1. Particle 0 is exact zero (base model)
|
|
decoded_p0 = transform.decode(Z[0:1]).squeeze(0)
|
|
assert torch.allclose(decoded_p0, base_vec, atol=1e-6), f"Particle 0 for dim={d} must match exact base vector"
|
|
|
|
# 2. For even N=30: particle N-1 (index 29) is also exact zero
|
|
decoded_plast = transform.decode(Z[29:30]).squeeze(0)
|
|
assert torch.allclose(decoded_plast, base_vec, atol=1e-6), f"Particle N-1 for dim={d} must match exact zero"
|
|
|
|
# 3. Antithetic pairs (particles 1..28 in 14 exact pairs)
|
|
z1 = Z[1:2]
|
|
z2 = Z[2:3]
|
|
assert torch.allclose(z1 + z2, torch.zeros_like(z1), atol=1e-6), "Antithetic pair latent sum must be zero"
|
|
|
|
delta1 = transform.decode(z1).squeeze(0) - transform.base_vec
|
|
delta2 = transform.decode(z2).squeeze(0) - transform.base_vec
|
|
assert torch.allclose(delta1 + delta2, torch.zeros_like(delta1), atol=1e-5), "Antithetic pair delta sum must be zero"
|
|
|
|
# 4. Latent swarm centroid is strictly zero
|
|
centroid_z = Z.mean(dim=0)
|
|
assert torch.allclose(centroid_z, torch.zeros_like(centroid_z), atol=1e-6), "Overall swarm centroid must be zero"
|
|
|
|
|
|
def test_transition_reevaluation_and_exact_accounting():
|
|
"""Verify objective size transition re-evaluates all pbests, rebuilds gbest, and asserts exact query/sample counts."""
|
|
device = torch.device("cpu")
|
|
base_model = make_compact_cnn(seed=41).to(device)
|
|
|
|
N_samples = 100
|
|
x_synth = torch.randn(N_samples, 1, 28, 28)
|
|
y_synth = torch.randint(0, 10, (N_samples,))
|
|
|
|
nested_subsets = {
|
|
20: torch.arange(20, dtype=torch.long),
|
|
50: torch.arange(50, dtype=torch.long),
|
|
}
|
|
|
|
transform = LatentTransform(base_model, latent_dim=290, device=device)
|
|
|
|
# Run 2-stage PSO: stage 0 (20 samples, 2 epochs), stage 1 (50 samples, 2 epochs), swarm_size = 10
|
|
# Stage 0: 2 * 10 = 20 queries, 20 * 20 = 400 sample evals. No transition re-eval.
|
|
# Stage 1 transition: 1 * 10 = 10 queries, 10 * 50 = 500 sample evals. Transition count = 10.
|
|
# Stage 1: 2 * 10 = 20 queries, 20 * 50 = 1000 sample evals.
|
|
# Total queries = 20 + 10 + 20 = 50.
|
|
# Total sample evals = 400 + 500 + 1000 = 1900.
|
|
res = run_latent_pso(
|
|
transform=transform,
|
|
base_model=base_model,
|
|
x_search=x_synth,
|
|
y_search=y_synth,
|
|
nested_subsets=nested_subsets,
|
|
schedule_str="20:2,50:2",
|
|
epochs=4,
|
|
swarm_size=10,
|
|
seed=42,
|
|
device=device,
|
|
)
|
|
|
|
assert res["transition_reevaluation_counts"] == 10, f"Expected 10 transition re-evaluations; got {res['transition_reevaluation_counts']}"
|
|
assert res["total_queries"] == 50, f"Expected 50 total queries; got {res['total_queries']}"
|
|
assert res["total_sample_evaluations"] == 1900, f"Expected 1900 sample evaluations; got {res['total_sample_evaluations']}"
|
|
assert len(res["stage_histories"]) == 4
|
|
|
|
|
|
def test_validation_only_elite_and_ensemble_selection():
|
|
"""Verify validation metrics correctly rank candidates and metric routines compute expected values."""
|
|
N = 100
|
|
C = 10
|
|
y_val = torch.randint(0, C, (N,))
|
|
|
|
# Candidate 1: Perfect predictions
|
|
probs_perfect = torch.zeros((N, C), dtype=torch.float32)
|
|
probs_perfect[torch.arange(N), y_val] = 1.0
|
|
|
|
# Candidate 2: Random noise
|
|
probs_random = torch.full((N, C), 1.0 / C, dtype=torch.float32)
|
|
|
|
m1 = evaluate_probabilistic_metrics(probs_perfect, y_val)
|
|
m2 = evaluate_probabilistic_metrics(probs_random, y_val)
|
|
|
|
assert m1["accuracy"] == 100.0
|
|
assert m1["nll"] < m2["nll"]
|
|
assert m1["brier"] < m2["brier"]
|
|
assert m1["ece"] <= 0.01
|
|
|
|
candidates = [
|
|
{"id": "cand2", "val_loss": m2["nll"], "val_acc": m2["accuracy"]},
|
|
{"id": "cand1", "val_loss": m1["nll"], "val_acc": m1["accuracy"]},
|
|
]
|
|
candidates.sort(key=lambda c: (c["val_loss"], -c["val_acc"]))
|
|
assert candidates[0]["id"] == "cand1"
|
|
|
|
diverse_candidates = [
|
|
{
|
|
"seed": 7,
|
|
"particle_idx": particle_idx,
|
|
"val_loss": 0.2 + 0.01 * particle_idx,
|
|
"val_acc": 90.0 - 0.25 * particle_idx,
|
|
"val_probs": torch.roll(probs_perfect, shifts=particle_idx, dims=1),
|
|
"latent_z": torch.full((4,), float(particle_idx)),
|
|
}
|
|
for particle_idx in range(3)
|
|
]
|
|
selected = select_diverse_candidates(
|
|
diverse_candidates, max_size=3, accuracy_window=2.0
|
|
)
|
|
assert len(selected) == 3
|
|
assert len({
|
|
(candidate["seed"], candidate["particle_idx"])
|
|
for candidate in selected
|
|
}) == 3
|
|
|
|
|
|
def test_cli_argument_validation():
|
|
"""Verify CLI argument validation rejects invalid parameters and accepts valid settings."""
|
|
parser = build_parser()
|
|
|
|
# Valid args
|
|
valid_args = parser.parse_args([
|
|
"--pilot-epochs", "160",
|
|
"--confirmation-epochs", "600",
|
|
"--confirmation-schedule", "2000:420,10000:135,50000:45",
|
|
"--seeds", "101", "102", "103",
|
|
"--dimensions", "290", "1024", "4096", "full"
|
|
])
|
|
validate_cli_args(valid_args)
|
|
|
|
# Invalid schedule sum mismatch
|
|
invalid_schedule = parser.parse_args([
|
|
"--confirmation-epochs", "600",
|
|
"--confirmation-schedule", "2000:400,10000:100,50000:50" # Sums to 550 != 600
|
|
])
|
|
with pytest.raises(ValueError, match="Schedule epoch sum"):
|
|
validate_cli_args(invalid_schedule)
|
|
|
|
# Invalid negative seed
|
|
invalid_seed = parser.parse_args(["--seeds", "-1"])
|
|
with pytest.raises(ValueError, match="Seeds must be non-negative"):
|
|
validate_cli_args(invalid_seed)
|
|
|
|
# Invalid duplicate seed
|
|
duplicate_seed = parser.parse_args(["--seeds", "101", "101"])
|
|
with pytest.raises(ValueError, match="Confirmation seeds must be unique"):
|
|
validate_cli_args(duplicate_seed)
|