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
This commit is contained in:
2026-09-07 22:03:25 +09:00
parent 9745bb7ad4
commit 813433000a
120 changed files with 198148 additions and 4745 deletions
+94 -71
View File
@@ -1,88 +1,111 @@
# %%
import os
import sys
import argparse
import torch
import torch.nn as nn
from pso import optimizer
import tensorflow as tf
from keras.datasets import mnist
from keras.layers import Conv2D, Dense, Dropout, Flatten, MaxPooling2D
from keras.models import Sequential
from pso import Optimizer
from cli import add_pso_args, build_optimizer_kwargs
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
def get_data(seed: int = 42):
from sklearn.decomposition import PCA
from torchvision.datasets import MNIST
train_dataset = MNIST(root="./data", train=True, download=True)
test_dataset = MNIST(root="./data", train=False, download=True)
def get_data():
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train_raw = (train_dataset.data[:3000].float() / 255.0).reshape(3000, -1).numpy()
y_train = train_dataset.targets[:3000].long()
x_train, x_test = x_train / 255.0, x_test / 255.0
x_train = x_train.reshape((60000, 28, 28, 1))
x_test = x_test.reshape((10000, 28, 28, 1))
x_test_raw = (test_dataset.data[:1000].float() / 255.0).reshape(1000, -1).numpy()
y_test = test_dataset.targets[:1000].long()
y_train, y_test = tf.one_hot(y_train, 10), tf.one_hot(y_test, 10)
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, x_test = tf.convert_to_tensor(x_train), tf.convert_to_tensor(x_test)
y_train, y_test = tf.convert_to_tensor(y_train), tf.convert_to_tensor(y_test)
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[0].shape} | y_train : {y_train[0].shape}")
print(f"x_test : {x_test[0].shape} | y_test : {y_test[0].shape}")
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():
model = Sequential()
model.add(
Conv2D(32, kernel_size=(5, 5), activation="relu", input_shape=(28, 28, 1))
def make_model(seed: int = 42):
torch.manual_seed(seed)
return nn.Linear(32, 10)
def main():
parser = argparse.ArgumentParser(description="PSO MNIST Benchmark Script")
add_pso_args(
parser,
defaults={
"method": "inertia",
"initialization": "model_noise",
"evaluation": "fixed_subset",
"convergence": "none",
"refinement": "adam",
"n_particles": 30,
"c0": None,
"c1": None,
"w_min": None,
"w_max": None,
"negative_swarm": 0.0,
"mutation_swarm": 0.02,
"particle_min": -3.0,
"particle_max": 3.0,
"velocity_limit_ratio": 0.025,
"boundary_strategy": "reflect",
"initial_position_noise": 0.05,
"seed": 42,
"epochs": 80,
"batch_size": 1000,
"fitness_size": 2000,
"renewal": "loss",
"output_dir": "output/mnist",
"checkpoint_interval": 25,
"refinement_epochs": 100,
"refinement_lr": 0.01,
},
)
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.5))
model.add(Conv2D(64, kernel_size=(3, 3), activation="relu"))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dropout(0.5))
model.add(Dense(256, activation="relu"))
model.add(Dense(128, activation="relu"))
model.add(Dense(10, activation="softmax"))
args = parser.parse_args()
return model
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": 1.49618, "c1": 1.49618, "w_min": 0.7298, "w_max": 0.7298},
)
pso_mnist = Optimizer(**kwargs)
print(f"Optimizer device: {pso_mnist.device}")
best_score = pso_mnist.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}")
# %%
model = make_model()
x_train, y_train, x_test, y_test = get_data()
pso_mnist = optimizer(
model,
loss="categorical_crossentropy",
n_particles=200,
c0=0.7,
c1=0.4,
w_min=0.1,
w_max=0.9,
negative_swarm=0.0,
mutation_swarm=0.05,
convergence_reset=True,
convergence_reset_patience=10,
convergence_reset_monitor="loss",
convergence_reset_min_delta=0.005,
)
best_score = pso_mnist.fit(
x_train,
y_train,
epochs=1000,
save_info=True,
log=2,
log_name="mnist",
renewal="loss",
check_point=25,
batch_size=5000,
validate_data=(x_test, y_test),
)
print("Done!")
sys.exit(0)
if __name__ == "__main__":
main()