Upload ai/training/train.py with huggingface_hub

ai/training/train.py

@@ -0,0 +1,280 @@
+import gc
+import glob
+import os
+import random
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+# Add project root to path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+from ai.models.network_torch import TransformerCardNet
+
+
+class ChunkDataset(Dataset):
+    """Dataset for a single chunk of data."""
+
+    def __init__(self, states, policies, winners, scores, turns):
+        self.states = torch.as_tensor(states, dtype=torch.float32)
+        self.policies = torch.as_tensor(policies, dtype=torch.float32)
+        self.winners = torch.as_tensor(winners, dtype=torch.float32).view(-1, 1)
+        self.scores = torch.as_tensor(scores, dtype=torch.float32).view(-1, 1)
+        self.turns = torch.as_tensor(turns, dtype=torch.float32).view(-1, 1)
+
+    def __len__(self):
+        return len(self.states)
+
+    def __getitem__(self, idx):
+        return (self.states[idx], self.policies[idx], self.winners[idx], self.scores[idx], self.turns[idx])
+
+
+def get_data_from_files(files):
+    """Loads and concatenates data from a list of files."""
+    all_states, all_policies, all_winners = [], [], []
+    all_scores, all_turns = [], []
+
+    for f in files:
+        data = np.load(f)
+        all_states.append(data["states"])
+        all_policies.append(data["policies"])
+        all_winners.append(data["winners"])
+        all_scores.append(data["scores"])
+        all_turns.append(data["turns_left"])
+
+    return (
+        np.concatenate(all_states),
+        np.concatenate(all_policies),
+        np.concatenate(all_winners),
+        np.concatenate(all_scores),
+        np.concatenate(all_turns),
+    )
+
+
+def train(data_pattern, epochs=20, batch_size=16384, lr=0.001, resume_path=None):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Training on {device} with batch size {batch_size}")
+
+    # Locate all matching files
+    all_files = sorted(glob.glob(data_pattern))
+    if not all_files:
+        print(f"Error: No files found matching {data_pattern}")
+        return
+
+    print(f"Found {len(all_files)} data chunks.")
+
+    # Reserve ~10% of files for validation (exclusive)
+    # If only 1 file, use it for BOTH training and validation
+    if len(all_files) == 1:
+        train_files = val_files = all_files
+    else:
+        val_count = max(1, len(all_files) // 10)
+        # Ensure at least one training file exists
+        if val_count >= len(all_files):
+            val_count = 0
+
+        rand_gen = random.Random(42)
+        shuffled_files = all_files.copy()
+        rand_gen.shuffle(shuffled_files)
+
+        train_files = shuffled_files[:-val_count] if val_count > 0 else shuffled_files
+        val_files = shuffled_files[-val_count:] if val_count > 0 else shuffled_files
+
+    print(f"Split: {len(train_files)} training files, {len(val_files)} validation files.")
+
+    # Initialize model
+    model = TransformerCardNet().to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, "min", patience=3, factor=0.5)
+
+    # Loss functions
+    policy_loss_fn = nn.CrossEntropyLoss(reduction="none")
+    win_loss_fn = nn.BCELoss()
+    mse_loss_fn = nn.MSELoss()
+
+    best_val_loss = float("inf")
+    start_epoch = 0
+
+    # Resume Logic
+    if resume_path and os.path.exists(resume_path):
+        print(f" -> Resuming from checkpoint: {resume_path}")
+        checkpoint = torch.load(resume_path, map_location=device)
+        model.load_state_dict(checkpoint["model_state"])
+        if "optimizer_state" in checkpoint:
+            optimizer.load_state_dict(checkpoint["optimizer_state"])
+        if "scheduler_state" in checkpoint:
+            scheduler.load_state_dict(checkpoint["scheduler_state"])
+        start_epoch = checkpoint.get("epoch", -1) + 1
+        best_val_loss = checkpoint.get("val_loss", float("inf"))
+        print(f" -> Starting from epoch {start_epoch + 1}")
+
+    try:
+        for epoch in range(start_epoch, epochs):
+            model.train()
+            train_losses = {"poly": 0, "win": 0, "score": 0, "turn": 0}
+            correct_policy = 0
+            total_policy = 0
+
+            current_train_files = train_files.copy()
+            random.shuffle(current_train_files)
+
+            pbar = tqdm(current_train_files, desc=f"Epoch {epoch + 1}/{epochs}")
+            for f_idx, f_path in enumerate(pbar):
+                try:
+                    s, pol, w, sc, tn = get_data_from_files([f_path])
+                    chunk_dataset = ChunkDataset(s, pol, w, sc, tn)
+                    chunk_loader = DataLoader(chunk_dataset, batch_size=batch_size, shuffle=True)
+
+                    for states, target_p, target_w, target_s, target_t in chunk_loader:
+                        states = states.to(device)
+                        target_p = target_p.to(device)
+                        target_w = target_w.to(device)
+                        target_s = target_s.to(device)
+                        target_t = target_t.to(device)
+
+                        optimizer.zero_grad()
+
+                        # Forward pass
+                        p_soft, w_pred, s_pred, t_pred = model(states)
+
+                        # 1. Policy Loss (Weighted)
+                        # Avoid log(0) by using logits or CrossEntropy internally
+                        # But TransformerCardNet returns softmax, so we use NLL or manual CE
+                        # For simplicity, let's use the probabilities directly or re-logit
+                        p_loss_raw = -torch.sum(target_p * torch.log(p_soft + 1e-8), dim=1)
+
+                        # Weight non-pass actions higher
+                        target_argmax = torch.max(target_p, dim=1)[1]
+                        weights = torch.ones(states.size(0), device=device)
+                        weights[target_argmax != 0] = 5.0
+                        p_loss = (p_loss_raw * weights).mean()
+
+                        # 2. Value Losses
+                        loss_win = win_loss_fn(w_pred, target_w)
+                        loss_score = mse_loss_fn(s_pred, target_s)
+                        loss_turn = mse_loss_fn(t_pred, target_t)
+
+                        total_loss = p_loss + loss_win + loss_score + loss_turn
+                        total_loss.backward()
+                        optimizer.step()
+
+                        train_losses["poly"] += p_loss.item()
+                        train_losses["win"] += loss_win.item()
+                        train_losses["score"] += loss_score.item()
+                        train_losses["turn"] += loss_turn.item()
+
+                        _, pred_action = torch.max(p_soft, 1)
+                        _, target_action = torch.max(target_p, 1)
+                        correct_policy += (pred_action == target_action).sum().item()
+                        total_policy += states.size(0)
+
+                    pbar.set_postfix(
+                        {
+                            "acc": f"{100 * correct_policy / total_policy:.1f}%",
+                            "win": f"{train_losses['win'] / (f_idx + 1):.3f}",
+                        }
+                    )
+
+                    del chunk_dataset, chunk_loader, s, pol, w, sc, tn
+                    gc.collect()
+                except Exception as e:
+                    print(f"Error processing chunk {f_path}: {e}")
+                    continue
+
+            # Validation
+            model.eval()
+            val_losses = {"poly": 0, "win": 0, "score": 0, "turn": 0}
+            val_correct = 0
+            val_total = 0
+            num_batches_val = 0
+
+            print(f" [Epoch {epoch + 1}] Validating...")
+            with torch.no_grad():
+                for f_path in val_files:
+                    try:
+                        s, pol, w, sc, tn = get_data_from_files([f_path])
+                        val_chunk = ChunkDataset(s, pol, w, sc, tn)
+                        val_loader = DataLoader(val_chunk, batch_size=batch_size, shuffle=False)
+
+                        for states, target_p, target_w, target_s, target_t in val_loader:
+                            states = states.to(device)
+                            target_p = target_p.to(device)
+                            target_w = target_w.to(device)
+                            target_s = target_s.to(device)
+                            target_t = target_t.to(device)
+
+                            p_soft, w_pred, s_pred, t_pred = model(states)
+
+                            p_loss_raw = -torch.sum(target_p * torch.log(p_soft + 1e-8), dim=1)
+                            val_losses["poly"] += p_loss_raw.mean().item()
+                            val_losses["win"] += win_loss_fn(w_pred, target_w).item()
+                            val_losses["score"] += mse_loss_fn(s_pred, target_s).item()
+                            val_losses["turn"] += mse_loss_fn(t_pred, target_t).item()
+
+                            num_batches_val += 1
+                            _, pred_action = torch.max(p_soft, 1)
+                            _, target_action = torch.max(target_p, 1)
+                            val_correct += (pred_action == target_action).sum().item()
+                            val_total += states.size(0)
+
+                        del val_chunk, val_loader, s, pol, w, sc, tn
+                        gc.collect()
+                    except Exception as e:
+                        print(f"Error validating chunk {f_path}: {e}")
+
+            if val_total > 0:
+                avg_val_win = val_losses["win"] / num_batches_val
+                avg_val_total = (val_losses["poly"] + val_losses["win"] + val_losses["score"]) / num_batches_val
+
+                scheduler.step(avg_val_total)
+
+                print(f"Epoch {epoch + 1} | Val WinLoss: {avg_val_win:.4f} | Acc: {100 * val_correct / val_total:.1f}%")
+
+                checkpoint = {
+                    "model_state": model.state_dict(),
+                    "optimizer_state": optimizer.state_dict(),
+                    "val_loss": avg_val_total,
+                    "epoch": epoch,
+                }
+                torch.save(checkpoint, f"ai/models/transformer_epoch_{epoch + 1}.pt")
+
+                if avg_val_total < best_val_loss:
+                    best_val_loss = avg_val_total
+                    torch.save(checkpoint, "ai/models/transformer_best.pt")
+
+    except KeyboardInterrupt:
+        print("\nTraining interrupted by user. Saving current state to ai/models/alphanet_interrupted.pt...")
+        # Save complete state for resumption
+        interrupted_state = {
+            "model_state": model.state_dict(),
+            "optimizer_state": optimizer.state_dict(),
+            "scheduler_state": scheduler.state_dict() if "scheduler" in locals() else None,
+            "epoch": epoch if "epoch" in locals() else 0,
+            "val_loss": best_val_loss,
+        }
+        torch.save(interrupted_state, "ai/models/alphanet_interrupted.pt")
+
+    torch.save(model.state_dict(), "ai/models/alphanet_final.pt")
+    print(f"Training complete. Best Val Loss: {best_val_loss:.4f}")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data", type=str, default="ai/data/alphazero_nightly_chunk_*.npz")
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--batch-size", type=int, default=16384)
+    parser.add_argument("--lr", type=float, default=0.001)
+    parser.add_argument(
+        "--resume", type=str, help="Path to checkpoint to resume from (e.g. ai/models/alphanet_interrupted.pt)"
+    )
+    args = parser.parse_args()
+
+    if not os.path.exists("ai/models"):
+        os.makedirs("ai/models")
+
+    train(args.data, epochs=args.epochs, batch_size=args.batch_size, lr=args.lr, resume_path=args.resume)