mirror of
https://github.com/jung-geun/PSO.git
synced 2026-09-20 14:11:48 +09:00
jung-geun
813433000a
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
89 lines
3.2 KiB
Python
89 lines
3.2 KiB
Python
import collections
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class ParameterCodec:
|
|
"""
|
|
Private parameter encoder/decoder for flattening and reconstructing PyTorch nn.Module parameters.
|
|
"""
|
|
|
|
def __init__(self, model: nn.Module):
|
|
if not isinstance(model, nn.Module):
|
|
raise TypeError("model must be an instance of torch.nn.Module")
|
|
|
|
params = list(model.named_parameters())
|
|
if not params:
|
|
raise ValueError("model contains zero trainable parameters")
|
|
|
|
self.names: list[str] = []
|
|
self.shapes: list[tuple[int, ...]] = []
|
|
self.numels: list[int] = []
|
|
dtypes: set[torch.dtype] = set()
|
|
|
|
total_size = 0
|
|
for name, p in params:
|
|
if not p.dtype.is_floating_point:
|
|
raise ValueError(
|
|
f"Parameter '{name}' has non-floating dtype {p.dtype}. Only floating-point parameters are supported."
|
|
)
|
|
self.names.append(name)
|
|
self.shapes.append(tuple(p.shape))
|
|
n = p.numel()
|
|
self.numels.append(n)
|
|
total_size += n
|
|
dtypes.add(p.dtype)
|
|
|
|
if len(dtypes) > 1:
|
|
raise ValueError(f"Mixed parameter dtypes found in model: {dtypes}")
|
|
|
|
self.dtype: torch.dtype = next(iter(dtypes))
|
|
self.size: int = total_size
|
|
|
|
def encode(self, model: nn.Module) -> torch.Tensor:
|
|
"""
|
|
Flattens model parameters into a single detached 1D torch.Tensor.
|
|
"""
|
|
params = [p for _, p in model.named_parameters()]
|
|
return nn.utils.parameters_to_vector(params).detach()
|
|
|
|
def apply_vector(self, vector: torch.Tensor, model: nn.Module) -> None:
|
|
"""
|
|
Applies a validated 1D parameter vector to model.parameters() in-place without storage aliasing.
|
|
"""
|
|
if not isinstance(vector, torch.Tensor) or vector.ndim != 1:
|
|
raise ValueError("vector must be a 1D torch.Tensor")
|
|
if vector.numel() != self.size:
|
|
raise ValueError(
|
|
f"Vector size {vector.numel()} does not match codec size {self.size}"
|
|
)
|
|
if vector.dtype != self.dtype:
|
|
raise ValueError(
|
|
f"Vector dtype {vector.dtype} does not match codec dtype {self.dtype}"
|
|
)
|
|
|
|
params = [p for _, p in model.named_parameters()]
|
|
if params:
|
|
target_device = params[0].device
|
|
if vector.device != target_device:
|
|
vector = vector.to(target_device)
|
|
|
|
with torch.no_grad():
|
|
offset = 0
|
|
for p, shape, numel in zip(params, self.shapes, self.numels):
|
|
p.copy_(vector[offset : offset + numel].reshape(shape))
|
|
offset += numel
|
|
|
|
def to_state_dict(
|
|
self, vector: torch.Tensor, eval_model: nn.Module
|
|
) -> collections.OrderedDict[str, torch.Tensor]:
|
|
"""
|
|
Applies vector to eval_model and returns a defensive CPU-cloned ordered state dict including buffers.
|
|
"""
|
|
self.apply_vector(vector, eval_model)
|
|
state_dict = eval_model.state_dict()
|
|
res = collections.OrderedDict()
|
|
for k, v in state_dict.items():
|
|
res[k] = v.detach().cpu().clone()
|
|
return res
|