ml/vae/training.py - main - sususer/ColonyGEN

ml/vae/training.py
406 строк18 KB
Chekr
f11
03 июн 2026, 13:53
03 июн 2026, 13:53d260831
from __future__ import annotations

import json
import time
from dataclasses import asdict
from pathlib import Path
from typing import Any

import numpy as np

from .codec import WorldTensorCodec
from .dataset import MapDataset, build_condition_vector
from .model import AdamOptimizer, NumpyVAE, VAEConfig

class VAETrainer:
    def __init__(
        self,
        dataset_root: str | Path,
        output_dir: str | Path,
        batch_size: int = 16,
        epochs: int = 40,
        learning_rate: float = 1e-3,
        beta_max: float = 0.02,
        beta_warmup_epochs: int = 10,
        seed: int = 42,
        train_limit: int | None = None,
        val_limit: int | None = None,
    ):
        self.dataset_root = Path(dataset_root)
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.batch_size = batch_size
        self.epochs = epochs
        self.learning_rate = learning_rate
        self.beta_max = beta_max
        self.beta_warmup_epochs = max(1, beta_warmup_epochs)
        self.seed = seed
        self.train_dataset = MapDataset(self.dataset_root, "train", limit=train_limit)
        self.codec = self.train_dataset.codec
        if self.codec.config.map_width != 48 or self.codec.config.map_height != 48:
            raise ValueError(
                "Legacy dataset size detected. "
                f"Expected 48x48, got {self.codec.config.map_width}x{self.codec.config.map_height}. "
                "Regenerate the dataset before training the VAE."
            )
        self.val_dataset = MapDataset(self.dataset_root, "val", codec=self.codec, limit=val_limit)
        self.condition_dim = int(self.train_dataset.load_conditions().shape[1])
        self.config = VAEConfig(
            map_width=self.codec.config.map_width,
            map_height=self.codec.config.map_height,
            tile_feature_dim=self.codec.feature_dim,
            condition_dim=self.condition_dim,
            seed=seed,
        )
        self.model = NumpyVAE(
            self.config,
            feature_weights=self.codec.feature_weights,
            terrain_class_weights=self._terrain_class_weights(),
        )
        self.optimizer = AdamOptimizer(learning_rate=learning_rate)

    def _terrain_class_weights(self) -> np.ndarray:
        weights = np.ones(len(self.codec.terrain_types), dtype=np.float32)
        weights[self.codec.terrain_to_idx["water"]] = 1.75
        weights[self.codec.terrain_to_idx["beach"]] = 1.10
        weights[self.codec.terrain_to_idx["forest"]] = 1.05
        return weights

    def _log(self, message: str) -> None:
        print(message, flush=True)

    def _progress_interval(self, total_batches: int) -> int:
        return max(1, total_batches // 10)

    def train(self) -> dict[str, Any]:
        self._log(
            f"[VAE][TRAIN] loading datasets root={self.dataset_root} "
            f"batch_size={self.batch_size} epochs={self.epochs} seed={self.seed}"
        )
        train_data = self.train_dataset.load_array()
        train_conditions = self.train_dataset.load_conditions()
        val_data = self.val_dataset.load_array()
        val_conditions = self.val_dataset.load_conditions()
        rng = np.random.default_rng(self.seed)
        best_val_loss = float("inf")
        best_path = self.output_dir / "vae_best.npz"
        history: list[dict[str, float]] = []
        total_train_batches = max(1, (len(train_data) + self.batch_size - 1) // self.batch_size)
        total_val_batches = max(1, (len(val_data) + self.batch_size - 1) // self.batch_size)
        train_started = time.perf_counter()

        self._log(
            f"[VAE][TRAIN] datasets ready train_samples={len(train_data)} "
            f"val_samples={len(val_data)} train_batches={total_train_batches} "
            f"val_batches={total_val_batches} feature_dim={self.codec.feature_dim}"
        )

        for epoch in range(1, self.epochs + 1):
            epoch_started = time.perf_counter()
            beta = min(self.beta_max, self.beta_max * epoch / self.beta_warmup_epochs)
            permutation = rng.permutation(len(train_data))
            shuffled = train_data[permutation]
            shuffled_conditions = train_conditions[permutation]
            train_sums = {"loss": 0.0, "reconstruction_loss": 0.0, "kl_loss": 0.0}
            gradient_norms = []
            batch_count = 0
            progress_interval = self._progress_interval(total_train_batches)

            self._log(
                f"[VAE][EPOCH {epoch:03d}] started beta={beta:.5f} "
                f"train_batches={total_train_batches}"
            )

            for start in range(0, len(shuffled), self.batch_size):
                batch = shuffled[start:start + self.batch_size]
                batch_condition = shuffled_conditions[start:start + self.batch_size]
                metrics, grads = self.model.loss_and_gradients(
                    batch,
                    condition=batch_condition,
                    beta=beta,
                    seed=self.seed + epoch * 1000 + batch_count,
                )
                grad_norm = self.optimizer.step(self.model.params, grads)
                gradient_norms.append(grad_norm)
                for key in train_sums:
                    train_sums[key] += metrics[key]
                batch_count += 1

                if (
                    batch_count == 1
                    or batch_count == total_train_batches
                    or batch_count % progress_interval == 0
                ):
                    self._log(
                        f"[VAE][EPOCH {epoch:03d}] train_batch={batch_count}/{total_train_batches} "
                        f"avg_loss={train_sums['loss'] / batch_count:.5f} "
                        f"last_loss={metrics['loss']:.5f} grad={grad_norm:.5f} "
                        f"elapsed={time.perf_counter() - epoch_started:.1f}s"
                    )

            train_metrics = {
                key: value / max(1, batch_count)
                for key, value in train_sums.items()
            }
            val_metrics = self.evaluate(val_data, val_conditions, beta=beta, phase_label=f"EPOCH {epoch:03d} VAL")
            epoch_metrics = {
                "epoch": float(epoch),
                "beta": float(beta),
                "gradient_norm": float(np.mean(gradient_norms) if gradient_norms else 0.0),
                "train_loss": train_metrics["loss"],
                "train_reconstruction_loss": train_metrics["reconstruction_loss"],
                "train_kl_loss": train_metrics["kl_loss"],
                "val_loss": val_metrics["loss"],
                "val_reconstruction_loss": val_metrics["reconstruction_loss"],
                "val_kl_loss": val_metrics["kl_loss"],
                "terrain_accuracy": val_metrics["terrain_accuracy"],
                "biome_accuracy": val_metrics["biome_accuracy"],
                "base_accuracy": val_metrics["base_accuracy"],
                "resource_mae": val_metrics["resource_mae"],
            }
            history.append(epoch_metrics)
            self._log(
                f"[VAE][EPOCH {epoch:03d}] "
                f"train={train_metrics['loss']:.5f} "
                f"val={val_metrics['loss']:.5f} "
                f"resource_mae={val_metrics['resource_mae']:.4f} "
                f"terrain={val_metrics['terrain_accuracy']:.3f} "
                f"biome={val_metrics['biome_accuracy']:.3f} "
                f"elapsed={time.perf_counter() - epoch_started:.1f}s"
            )
            if val_metrics["loss"] < best_val_loss:
                best_val_loss = val_metrics["loss"]
                self.model.save(
                    best_path,
                    codec_config=self.codec.to_dict(),
                    extra_metadata={
                        "best_epoch": epoch,
                        "best_val_loss": best_val_loss,
                        "training_config": self._training_config_dict(),
                    },
                )
                self._log(
                    f"[VAE][EPOCH {epoch:03d}] new_best val_loss={best_val_loss:.5f} "
                    f"saved={best_path}"
                )

        history_path = self.output_dir / "history.json"
        history_path.write_text(json.dumps(history, indent=2), encoding="utf-8")
        summary = {
            "best_model": str(best_path),
            "history_path": str(history_path),
            "final_epoch": history[-1] if history else {},
        }
        summary_path = self.output_dir / "summary.json"
        summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
        self._log(
            f"[VAE][TRAIN] completed best_model={best_path} history={history_path} "
            f"summary={summary_path} elapsed={time.perf_counter() - train_started:.1f}s"
        )
        return summary

    def evaluate(
        self,
        data: np.ndarray,
        conditions: np.ndarray,
        beta: float,
        phase_label: str = "VAL",
    ) -> dict[str, float]:
        sums = {
            "loss": 0.0,
            "reconstruction_loss": 0.0,
            "kl_loss": 0.0,
            "terrain_accuracy": 0.0,
            "biome_accuracy": 0.0,
            "base_accuracy": 0.0,
            "resource_mae": 0.0,
        }
        total_batches = max(1, (len(data) + self.batch_size - 1) // self.batch_size)
        progress_interval = self._progress_interval(total_batches)
        phase_started = time.perf_counter()

        self._log(
            f"[VAE][{phase_label}] started samples={len(data)} batches={total_batches} beta={beta:.5f}"
        )

        batch_count = 0
        for start in range(0, len(data), self.batch_size):
            batch = data[start:start + self.batch_size]
            batch_condition = conditions[start:start + self.batch_size]
            metrics = self.model.evaluate_batch(batch, condition=batch_condition, beta=beta)
            reconstructed = self.model.reconstruct(batch, condition=batch_condition)
            recon_metrics = self._batch_reconstruction_metrics(batch, reconstructed)
            for key in ("loss", "reconstruction_loss", "kl_loss"):
                sums[key] += metrics[key]
            for key in ("terrain_accuracy", "biome_accuracy", "base_accuracy", "resource_mae"):
                sums[key] += recon_metrics[key]
            batch_count += 1

            if (
                batch_count == 1
                or batch_count == total_batches
                or batch_count % progress_interval == 0
            ):
                self._log(
                    f"[VAE][{phase_label}] batch={batch_count}/{total_batches} "
                    f"loss={metrics['loss']:.5f} terrain={recon_metrics['terrain_accuracy']:.3f} "
                    f"biome={recon_metrics['biome_accuracy']:.3f} "
                    f"resource_mae={recon_metrics['resource_mae']:.4f} "
                    f"elapsed={time.perf_counter() - phase_started:.1f}s"
                )

        result = {
            key: value / max(1, batch_count)
            for key, value in sums.items()
        }

        self._log(
            f"[VAE][{phase_label}] completed loss={result['loss']:.5f} "
            f"terrain={result['terrain_accuracy']:.3f} biome={result['biome_accuracy']:.3f} "
            f"resource_mae={result['resource_mae']:.4f} elapsed={time.perf_counter() - phase_started:.1f}s"
        )
        return result

    def sample(self, model_path: str | Path, output_dir: str | Path, count: int, temperature: float, seed: int) -> list[str]:
        self._log(
            f"[VAE][SAMPLE] loading model={model_path} dataset_root={self.dataset_root} "
            f"count={count} temperature={temperature:.2f} seed={seed}"
        )
        model, codec_config, _ = NumpyVAE.load(model_path)
        codec = self.codec if not codec_config else WorldTensorCodec.from_manifest(codec_config)
        target_dir = Path(output_dir)
        target_dir.mkdir(parents=True, exist_ok=True)
        tensors = model.sample(count=count, temperature=temperature, seed=seed)
        written = []
        progress_interval = self._progress_interval(max(1, len(tensors)))
        for idx, tensor in enumerate(tensors):
            payload = codec.decode_tensor(
                tensor,
                metadata={
                    "generator": "numpy_vae",
                    "temperature": temperature,
                    "sample_index": idx,
                },
                split="generated",
            )
            path = target_dir / f"sample_{idx:04d}.json"
            path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
            written.append(str(path))

            current = idx + 1
            if current == 1 or current == len(tensors) or current % progress_interval == 0:
                self._log(f"[VAE][SAMPLE] written={current}/{len(tensors)} file={path.name}")

        self._log(f"[VAE][SAMPLE] completed output_dir={target_dir} written={len(written)}")
        return written

    def reconstruct_files(self, model_path: str | Path, output_dir: str | Path, split: str, count: int) -> list[str]:
        self._log(
            f"[VAE][RECON] loading model={model_path} dataset_root={self.dataset_root} "
            f"split={split} count={count}"
        )
        model, codec_config, _ = NumpyVAE.load(model_path)
        dataset = MapDataset(self.dataset_root, split, codec=self.codec if not codec_config else WorldTensorCodec.from_manifest(codec_config), limit=count)
        tensors = dataset.load_array()
        conditions = dataset.load_conditions()
        recon = model.reconstruct(tensors, condition=conditions)
        target_dir = Path(output_dir)
        target_dir.mkdir(parents=True, exist_ok=True)
        written = []
        progress_interval = self._progress_interval(max(1, len(recon)))
        for idx, tensor in enumerate(recon):
            payload = dataset.codec.decode_tensor(
                tensor,
                metadata={
                    "generator": "numpy_vae_reconstruction",
                    "source_split": split,
                    "source_file": dataset.paths[idx].name,
                },
                split=f"reconstructed_{split}",
            )
            path = target_dir / f"reconstruction_{idx:04d}.json"
            path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
            written.append(str(path))

            current = idx + 1
            if current == 1 or current == len(recon) or current % progress_interval == 0:
                self._log(f"[VAE][RECON] written={current}/{len(recon)} file={path.name}")

        self._log(f"[VAE][RECON] completed output_dir={target_dir} written={len(written)}")
        return written

    def _batch_reconstruction_metrics(self, original: np.ndarray, reconstructed: np.ndarray) -> dict[str, float]:
        metrics = {
            "terrain_accuracy": 0.0,
            "biome_accuracy": 0.0,
            "base_accuracy": 0.0,
            "resource_mae": 0.0,
        }
        for src, dst in zip(original, reconstructed, strict=False):
            sample_metrics = self.codec.reconstruction_metrics(src, dst)
            for key, value in sample_metrics.items():
                metrics[key] += value
        return {
            key: value / max(1, len(original))
            for key, value in metrics.items()
        }

    def _training_config_dict(self) -> dict[str, Any]:
        return {
            "batch_size": self.batch_size,
            "epochs": self.epochs,
            "learning_rate": self.learning_rate,
            "beta_max": self.beta_max,
            "beta_warmup_epochs": self.beta_warmup_epochs,
            "seed": self.seed,
            "model": asdict(self.config),
            "terrain_class_weights": self.model.terrain_class_weights.tolist(),
        }

    def sample_conditioned(
        self,
        model_path: str | Path,
        output_dir: str | Path,
        count: int,
        temperature: float,
        seed: int,
        difficulty_score: float,
        entropy_score: float,
    ) -> list[str]:
        self._log(
            f"[VAE][SAMPLE][COND] loading model={model_path} dataset_root={self.dataset_root} "
            f"count={count} temperature={temperature:.2f} seed={seed} "
            f"difficulty={difficulty_score:.2f} entropy={entropy_score:.3f}"
        )
        model, codec_config, _ = NumpyVAE.load(model_path)
        codec = self.codec if not codec_config else WorldTensorCodec.from_manifest(codec_config)
        target_dir = Path(output_dir)
        target_dir.mkdir(parents=True, exist_ok=True)
        condition = build_condition_vector(difficulty_score, entropy_score)
        tensors = model.sample(count=count, temperature=temperature, seed=seed, condition=condition)
        written = []
        progress_interval = self._progress_interval(max(1, len(tensors)))
        for idx, tensor in enumerate(tensors):
            payload = codec.decode_tensor(
                tensor,
                metadata={
                    "generator": "numpy_cvae",
                    "temperature": temperature,
                    "sample_index": idx,
                    "target_condition": {
                        "difficulty_score": float(difficulty_score),
                        "entropy_score": float(entropy_score),
                    },
                },
                split="generated",
            )
            path = target_dir / f"sample_{idx:04d}.json"
            path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
            written.append(str(path))

            current = idx + 1
            if current == 1 or current == len(tensors) or current % progress_interval == 0:
                self._log(f"[VAE][SAMPLE][COND] written={current}/{len(tensors)} file={path.name}")

        self._log(f"[VAE][SAMPLE][COND] completed output_dir={target_dir} written={len(written)}")
        return written