Files
PSO/test/fashion_mnist.py
T
jung-geun 813433000a feat: modernize PSO and add convergence research
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
2026-09-07 22:03:25 +09:00

111 lines
3.3 KiB
Python

import argparse
import torch
import torch.nn as nn
from pso import Optimizer
from cli import add_pso_args, build_optimizer_kwargs
def get_data(seed: int = 42):
from sklearn.decomposition import PCA
from torchvision.datasets import FashionMNIST
train_dataset = FashionMNIST(root="./data", train=True, download=True)
test_dataset = FashionMNIST(root="./data", train=False, download=True)
x_train_raw = (train_dataset.data[:3000].float() / 255.0).reshape(3000, -1).numpy()
y_train = train_dataset.targets[:3000].long()
x_test_raw = (test_dataset.data[:1000].float() / 255.0).reshape(1000, -1).numpy()
y_test = test_dataset.targets[:1000].long()
pca = PCA(n_components=32, whiten=True, random_state=seed)
x_train_pca = pca.fit_transform(x_train_raw)
x_test_pca = pca.transform(x_test_raw)
x_train = torch.tensor(x_train_pca, dtype=torch.float32)
x_test = torch.tensor(x_test_pca, dtype=torch.float32)
print(f"x_train : {x_train.shape} | y_train : {y_train.shape}")
print(f"x_test : {x_test.shape} | y_test : {y_test.shape}")
return x_train, y_train, x_test, y_test
def make_model(seed: int = 42):
torch.manual_seed(seed)
return nn.Linear(32, 10)
def main():
parser = argparse.ArgumentParser(description="PSO Fashion-MNIST Benchmark Script")
add_pso_args(
parser,
defaults={
"method": "original",
"initialization": "model_noise",
"evaluation": "fixed_subset",
"convergence": "particle_reset",
"refinement": "adam",
"n_particles": 30,
"c0": None,
"c1": None,
"w_min": None,
"w_max": None,
"negative_swarm": 0.0,
"mutation_swarm": 0.05,
"particle_min": -3.0,
"particle_max": 3.0,
"velocity_limit_ratio": 0.1,
"boundary_strategy": "reflect",
"seed": 42,
"epochs": 80,
"batch_size": 1000,
"fitness_size": 2000,
"renewal": "loss",
"output_dir": "output/fashion_mnist",
"checkpoint_interval": 25,
"refinement_epochs": 10,
"refinement_lr": 0.001,
},
)
args = parser.parse_args()
model = make_model(seed=args.seed)
x_train, y_train, x_test, y_test = get_data(seed=args.seed)
fitness_size = args.fitness_size if args.evaluation == "fixed_subset" else None
refinement_epochs = args.refinement_epochs if args.refinement == "adam" else 0
kwargs = build_optimizer_kwargs(
args,
model=model,
loss=nn.CrossEntropyLoss(),
task="multiclass",
inertia_profile={"c0": 0.7, "c1": 0.5, "w_min": 0.1, "w_max": 0.8},
)
pso_fashion = Optimizer(**kwargs)
print(f"Optimizer device: {pso_fashion.device}")
best_score = pso_fashion.fit(
x_train,
y_train,
epochs=args.epochs,
batch_size=args.batch_size,
fitness_size=fitness_size,
renewal=args.renewal,
validation_data=(x_test, y_test),
output_dir=args.output_dir,
checkpoint_interval=25,
save_info=True,
refinement_epochs=refinement_epochs,
refinement_lr=args.refinement_lr,
)
print(f"Done! Best score: {best_score}")
if __name__ == "__main__":
main()