grok lock in pt 2

training.py

@@ -1,6 +1,3 @@
-# ===============================
-# IMPORTS
-# ===============================
 import os
 import json
 import math
@@ -10,35 +7,31 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
+
 from torch.utils.data import Dataset, DataLoader
 from sklearn.model_selection import train_test_split
-from sklearn.preprocessing import StandardScaler
+from sklearn.preprocessing import LabelEncoder, StandardScaler
 from multiprocessing import Pool, cpu_count
 from functools import partial
 from tqdm import tqdm
 
-# ===============================
-# DEVICE
-# ===============================
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
-if device.type == "cuda":
-    print("GPU:", torch.cuda.get_device_name(0))
-    torch.backends.cudnn.benchmark = True
 
-# ===============================
-# DATA LOADING
-# ===============================
 def load_kaggle_asl_data(base_path):
     train_df = pd.read_csv(os.path.join(base_path, "train.csv"))
     with open(os.path.join(base_path, "sign_to_prediction_index_map.json")) as f:
         sign_to_idx = json.load(f)
     return train_df, sign_to_idx
 
+
 def extract_hand_landmarks_from_parquet(path):
     try:
         df = pd.read_parquet(path)
-        hand = df[df["type"].isin(["left_hand", "right_hand"])]
+        # Take either left or right hand - prefer the one with more landmarks
+        left = df[df["type"] == "left_hand"]
+        right = df[df["type"] == "right_hand"]
+
+        hand = left if len(left) >= len(right) else right
+
         if len(hand) == 0:
             return None
 
@@ -53,115 +46,62 @@ def extract_hand_landmarks_from_parquet(path):
                 if len(lm) == 0:
                     lm_list.append([0.0, 0.0, 0.0])
                 else:
-                    lm_list.append([lm['x'].values[0], lm['y'].values[0], lm['z'].values[0]])
+                    lm_list.append([
+                        float(lm['x'].iloc[0]),
+                        float(lm['y'].iloc[0]),
+                        float(lm['z'].iloc[0])
+                    ])
             landmarks_seq.append(lm_list)
 
         return np.array(landmarks_seq, dtype=np.float32)  # (T, 21, 3)
-    except:
+    except Exception:
         return None
 
-def get_features_sequence(landmarks_seq, max_frames=96):
+
+def get_features_sequence(landmarks_seq, max_frames=100):
     if landmarks_seq is None or len(landmarks_seq) == 0:
         return None
 
-    # Center on wrist (landmark 0)
+    # Center on wrist
     landmarks_seq = landmarks_seq - landmarks_seq[:, 0:1, :]
 
-    # Rough scale normalization (using index finger length as reference)
+    # Better scale: distance between index finger tip and middle finger tip
-    scale = np.linalg.norm(landmarks_seq[:, 8] - landmarks_seq[:, 5], axis=1, keepdims=True)
+    scale = np.linalg.norm(landmarks_seq[:, 8] - landmarks_seq[:, 12], axis=1, keepdims=True)
     scale = np.maximum(scale, 1e-6)
-    landmarks_seq /= scale
+    landmarks_seq = landmarks_seq / scale
 
-    # Flatten → (T, 63)
+    # Flatten to (T, 63)
     seq = landmarks_seq.reshape(landmarks_seq.shape[0], -1)
 
-    # Pad / truncate
+    # Pad or truncate
     if len(seq) < max_frames:
         pad = np.zeros((max_frames - len(seq), seq.shape[1]), dtype=np.float32)
         seq = np.concatenate([seq, pad], axis=0)
     else:
         seq = seq[:max_frames]
 
-    return seq.astype(np.float32)
+    return seq
 
-def process_row(row, base_path, max_frames=96):
+
+def process_row(row, base_path, max_frames=100):
     path = os.path.join(base_path, row['path'])
     if not os.path.exists(path):
         return None, None
-    lm = extract_hand_landmarks_from_parquet(path)
+
-    feat = get_features_sequence(lm, max_frames)
+    try:
-    if feat is None:
+        lm_seq = extract_hand_landmarks_from_parquet(path)
+        if lm_seq is None:
             return None, None
-    return feat, row['sign']
 
-# ===============================
+        feat_seq = get_features_sequence(lm_seq, max_frames)
-# LOAD & PROCESS (with progress)
+        if feat_seq is None:
-# ===============================
+            return None, None
-base_path = "asl_kaggle"  # ← change if needed
-train_df, sign_to_idx = load_kaggle_asl_data(base_path)
 
-print("Processing videos...")
+        return feat_seq, row['sign']
-rows = [row for _, row in train_df.iterrows()]
+    except Exception:
+        return None, None
 
-with Pool(cpu_count()) as pool:
-    results = list(tqdm(pool.imap(
-        partial(process_row, base_path=base_path, max_frames=96),
-        rows
-    ), total=len(rows)))
 
-X, y = [], []
-for feat, sign in results:
-    if feat is not None:
-        X.append(feat)
-        y.append(sign)
-
-X = np.stack(X)  # (N, T, 63)
-print(f"Loaded {len(X)} valid samples  |  shape: {X.shape}")
-
-# Global normalization (very important!)
-scaler = StandardScaler()
-X_reshaped = X.reshape(-1, X.shape[-1])
-X_reshaped = scaler.fit_transform(X_reshaped)
-X = X_reshaped.reshape(X.shape)
-
-# ===============================
-# LABELS
-# ===============================
-from sklearn.preprocessing import LabelEncoder
-le = LabelEncoder()
-y = le.fit_transform(y)
-num_classes = len(le.classes_)
-print(f"Classes: {num_classes}")
-
-# ===============================
-# SPLIT
-# ===============================
-X_train, X_test, y_train, y_test = train_test_split(
-    X, y, test_size=0.15, stratify=y, random_state=42
-)
-
-# ===============================
-# DATASET + DATALOADER
-# ===============================
-class ASLSequenceDataset(Dataset):
-    def __init__(self, X, y):
-        self.X = torch.from_numpy(X).float()
-        self.y = torch.from_numpy(y).long()
-
-    def __len__(self):
-        return len(self.X)
-
-    def __getitem__(self, idx):
-        return self.X[idx], self.y[idx]
-
-train_loader = DataLoader(ASLSequenceDataset(X_train, y_train),
-                          batch_size=64, shuffle=True, num_workers=4, pin_memory=True)
-test_loader  = DataLoader(ASLSequenceDataset(X_test, y_test),
-                          batch_size=96, shuffle=False, num_workers=4, pin_memory=True)
-
-# ===============================
-# MODEL
-# ===============================
 class PositionalEncoding(nn.Module):
     def __init__(self, d_model, max_len=128):
         super().__init__()
@@ -175,12 +115,12 @@ class PositionalEncoding(nn.Module):
     def forward(self, x):
         return x + self.pe[:, :x.size(1)]
 
+
 class TransformerASL(nn.Module):
     def __init__(self, input_dim=63, num_classes=250, d_model=192, nhead=6, num_layers=4):
         super().__init__()
         self.proj = nn.Linear(input_dim, d_model)
         self.norm_in = nn.LayerNorm(d_model)
-
         self.pos = PositionalEncoding(d_model)
 
         encoder_layer = nn.TransformerEncoderLayer(
@@ -204,41 +144,150 @@ class TransformerASL(nn.Module):
         x = self.proj(x)
         x = self.norm_in(x)
         x = self.pos(x)
-
         x = self.encoder(x, src_key_padding_mask=key_padding_mask)
         x = x.mean(dim=1)  # global average pooling
-        x = self.head(x)
+        return self.head(x)
-        return x
+
+
+def create_padding_mask(lengths, max_len):
+    return torch.arange(max_len, device=lengths.device)[None, :] >= lengths[:, None]
+
+
+def main():
+    # ===============================
+    # DEVICE SETUP
+    # ===============================
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    if device.type == "cuda":
+        print("GPU:", torch.cuda.get_device_name(0))
+
+    # ===============================
+    # PATHS & PARAMETERS
+    # ===============================
+    base_path = "asl_kaggle"  # ← CHANGE THIS TO YOUR ACTUAL FOLDER
+    max_frames = 100
+
+    # ===============================
+    # DATA PROCESSING
+    # ===============================
+    print("Loading metadata...")
+    train_df, sign_to_idx = load_kaggle_asl_data(base_path)
+
+    print(f"Processing {len(train_df)} videos...")
+    rows = [row for _, row in train_df.iterrows()]
+
+    with Pool(cpu_count()) as pool:
+        results = list(tqdm(
+            pool.imap(
+                partial(process_row, base_path=base_path, max_frames=max_frames),
+                rows
+            ),
+            total=len(rows),
+            desc="Processing"
+        ))
+
+    X, y = [], []
+    for feat, sign in results:
+        if feat is not None:
+            X.append(feat)
+            y.append(sign)
+
+    if not X:
+        print("No valid sequences found!")
+        return
+
+    X = np.stack(X)
+    print(f"Loaded {len(X)} valid samples | shape: {X.shape}")
+
+    # Global normalization - very important!
+    print("Before global norm → mean:", X.mean(), "std:", X.std())
+    X = np.clip(X, -5.0, 5.0)  # prevent crazy outliers
+    mean = X.mean(axis=(0, 1), keepdims=True)
+    std = X.std(axis=(0, 1), keepdims=True) + 1e-8
+    X = (X - mean) / std
+    print("After global norm  → mean:", X.mean(), "std:", X.std())
+
+    # ===============================
+    # LABELS
+    # ===============================
+    le = LabelEncoder()
+    y = le.fit_transform(y)
+    num_classes = len(le.classes_)
+    print(f"Number of classes: {num_classes}")
+
+    # ===============================
+    # SPLIT
+    # ===============================
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y, test_size=0.15, stratify=y, random_state=42
+    )
+
+    # ===============================
+    # DATASET & LOADERS
+    # ===============================
+    class ASLSequenceDataset(Dataset):
+        def __init__(self, X, y):
+            self.X = torch.from_numpy(X).float()
+            self.y = torch.from_numpy(y).long()
+
+        def __len__(self):
+            return len(self.X)
+
+        def __getitem__(self, idx):
+            return self.X[idx], self.y[idx]
+
+    train_loader = DataLoader(
+        ASLSequenceDataset(X_train, y_train),
+        batch_size=64,
+        shuffle=True,
+        num_workers=4,
+        pin_memory=True
+    )
+
+    test_loader = DataLoader(
+        ASLSequenceDataset(X_test, y_test),
+        batch_size=96,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True
+    )
+
+    # ===============================
+    # MODEL
+    # ===============================
+    model = TransformerASL(
+        input_dim=63,
+        num_classes=num_classes,
+        d_model=192,
+        nhead=6,
+        num_layers=4
+    ).to(device)
 
-model = TransformerASL(input_dim=63, num_classes=num_classes).to(device)
     print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
 
     # ===============================
     # TRAINING SETUP
     # ===============================
     criterion = nn.CrossEntropyLoss(label_smoothing=0.05)
-optimizer = optim.AdamW(model.parameters(), lr=8e-4, weight_decay=1e-4, betas=(0.9, 0.98))
+    optimizer = optim.AdamW(model.parameters(), lr=5e-4, weight_decay=1e-4)
-scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=15, T_mult=2)
+    scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10)
 
     # ===============================
-# TRAIN / EVAL
+    # TRAIN / EVAL FUNCTIONS
     # ===============================
-def create_padding_mask(seq_len, max_len):
-    # True = ignore this position
-    return torch.arange(max_len, device=device)[None, :] >= seq_len[:, None]
-
     def train_epoch():
         model.train()
         total_loss = 0
         correct = 0
         total = 0
 
-    for x, y in tqdm(train_loader, desc="Train"):
+        for x, y in tqdm(train_loader, desc="Training"):
             x, y = x.to(device), y.to(device)
 
-        # Very simple length heuristic (can be improved later)
+            # Rough length estimation
-        real_lengths = (x.abs().sum(dim=2) > 1e-6).sum(dim=1)
+            lengths = (x.abs().sum(dim=2) > 1e-5).sum(dim=1)
-        mask = create_padding_mask(real_lengths, x.size(1))
+            mask = create_padding_mask(lengths, x.size(1))
 
             optimizer.zero_grad(set_to_none=True)
             logits = model(x, key_padding_mask=mask)
@@ -246,19 +295,17 @@ def train_epoch():
             loss = criterion(logits, y)
             loss.backward()
 
-        # STRONG clipping — very important for landmarks
             grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.8)
 
+            if torch.isnan(loss) or grad_norm > 20:
+                print(f"Warning - large grad or NaN! norm = {grad_norm:.2f}")
+
             optimizer.step()
 
             total_loss += loss.item()
             correct += (logits.argmax(dim=-1) == y).sum().item()
             total += y.size(0)
 
-        # Debug exploding gradients
-        if torch.isnan(loss) or grad_norm > 50:
-            print(f"WARNING - NaN or huge grad! norm={grad_norm:.2f}")
-
         return total_loss / len(train_loader), correct / total * 100
 
     @torch.no_grad()
@@ -268,27 +315,27 @@ def evaluate():
         total = 0
         for x, y in test_loader:
             x, y = x.to(device), y.to(device)
-        real_lengths = (x.abs().sum(dim=2) > 1e-6).sum(dim=1)
+            lengths = (x.abs().sum(dim=2) > 1e-5).sum(dim=1)
-        mask = create_padding_mask(real_lengths, x.size(1))
+            mask = create_padding_mask(lengths, x.size(1))
 
             logits = model(x, key_padding_mask=mask)
             correct += (logits.argmax(dim=-1) == y).sum().item()
             total += y.size(0)
-    return correct / total * 100
+        return correct / total * 100 if total > 0 else 0
 
     # ===============================
     # TRAINING LOOP
     # ===============================
     best_acc = 0
-patience = 18
+    patience = 15
     wait = 0
-epochs = 80
+    epochs = 60
 
     for epoch in range(epochs):
         loss, train_acc = train_epoch()
         test_acc = evaluate()
 
-    print(f"[{epoch+1:2d}/{epochs}]  loss: {loss:.4f}  |  train: {train_acc:.2f}%  |  test: {test_acc:.2f}%")
+        print(f"Epoch {epoch + 1:2d}/{epochs} | Loss: {loss:.4f} | Train: {train_acc:.2f}% | Test: {test_acc:.2f}%")
 
         scheduler.step()
 
@@ -298,14 +345,19 @@ for epoch in range(epochs):
             torch.save({
                 'model': model.state_dict(),
                 'optimizer': optimizer.state_dict(),
-            'scaler': scaler,
+                'label_encoder': le.classes_,
-            'label_encoder_classes': le.classes_
+                'epoch': epoch,
+                'acc': best_acc
             }, "best_asl_transformer.pth")
-        print("→ Saved new best model")
+            print("  → New best model saved")
         else:
             wait += 1
             if wait >= patience:
                 print("Early stopping triggered")
                 break
 
-print(f"\nBest test accuracy achieved: {best_acc:.2f}%")
+    print(f"\nTraining finished. Best test accuracy: {best_acc:.2f}%")
+
+
+if __name__ == '__main__':
+    main()