From ab907280ddedca702bd018ffee38d657465bda63 Mon Sep 17 00:00:00 2001 From: Stupdi Go Date: Sat, 10 Jan 2026 23:38:17 -0600 Subject: [PATCH] chatgpt lock tf in --- training.py | 250 +++++++++++++++++++--------------------------------- 1 file changed, 91 insertions(+), 159 deletions(-) diff --git a/training.py b/training.py index d1dfe04..0f2c267 100644 --- a/training.py +++ b/training.py @@ -10,34 +10,31 @@ import torch.optim as optim from torch.utils.data import Dataset, DataLoader from sklearn.model_selection import train_test_split -from sklearn.preprocessing import LabelEncoder, StandardScaler +from sklearn.preprocessing import LabelEncoder from multiprocessing import Pool, cpu_count from functools import partial from tqdm import tqdm from collections import Counter - +# =============================== +# 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) 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 - frames = sorted(hand['frame'].unique()) landmarks_seq = [] - for frame in frames: lm_frame = hand[hand['frame'] == frame] lm_list = [] @@ -52,12 +49,10 @@ def extract_hand_landmarks_from_parquet(path): float(lm['z'].iloc[0]) ]) landmarks_seq.append(lm_list) - return np.array(landmarks_seq, dtype=np.float32) except: return None - def get_features_sequence(landmarks_seq, max_frames=100): if landmarks_seq is None or len(landmarks_seq) == 0: return None @@ -65,45 +60,51 @@ def get_features_sequence(landmarks_seq, max_frames=100): # Center on wrist landmarks_seq -= landmarks_seq[:, 0:1, :] - # Scale using index → middle finger tip distance (more stable than single point) - scale = np.linalg.norm(landmarks_seq[:, 8] - landmarks_seq[:, 12], axis=1, keepdims=True) + # Robust scale: wrist → middle finger MCP + scale = np.linalg.norm(landmarks_seq[:,0] - landmarks_seq[:,9], axis=1, keepdims=True) scale = np.maximum(scale, 1e-6) - landmarks_seq /= scale + landmarks_seq /= scale[:, np.newaxis, :] # Flatten seq = landmarks_seq.reshape(landmarks_seq.shape[0], -1) # Pad / truncate - if len(seq) < max_frames: - pad = np.zeros((max_frames - len(seq), seq.shape[1]), dtype=np.float32) + T = seq.shape[0] + if T < max_frames: + pad = np.zeros((max_frames - T, seq.shape[1]), dtype=np.float32) seq = np.concatenate([seq, pad]) else: seq = seq[:max_frames] - return seq + # Mask for valid frames + valid_mask = np.zeros(max_frames, dtype=bool) + valid_mask[:min(T, max_frames)] = True + return seq, valid_mask 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 + return None, None, None try: lm = extract_hand_landmarks_from_parquet(path) if lm is None: - return None, None - feat = get_features_sequence(lm, max_frames) + return None, None, None + feat, mask = get_features_sequence(lm, max_frames) if feat is None: - return None, None - return feat, row["sign"] + return None, None, None + return feat, mask, row["sign"] except: - return None, None - + return None, None, None +# =============================== +# TRANSFORMER MODEL +# =============================== class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=128): super().__init__() pe = torch.zeros(max_len, d_model) - position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) + position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) @@ -112,28 +113,25 @@ 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): + def __init__(self, input_dim=63, num_classes=250, d_model=128, nhead=4, num_layers=2): super().__init__() self.proj = nn.Linear(input_dim, d_model) self.norm_in = nn.LayerNorm(d_model) self.pos = PositionalEncoding(d_model) - enc_layer = nn.TransformerEncoderLayer( d_model=d_model, nhead=nhead, - dim_feedforward=d_model * 4, - dropout=0.15, + dim_feedforward=d_model*4, + dropout=0.1, activation='gelu', batch_first=True, norm_first=True ) self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers) - self.head = nn.Sequential( nn.LayerNorm(d_model), - nn.Dropout(0.25), + nn.Dropout(0.2), nn.Linear(d_model, num_classes) ) @@ -142,209 +140,144 @@ class TransformerASL(nn.Module): 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 = x.mean(dim=1) return self.head(x) +def create_padding_mask(valid_masks): + # valid_masks: (B,T) bool, True for valid + return ~valid_masks # True in mask = positions to ignore -def create_padding_mask(lengths, max_len): - return torch.arange(max_len, device=lengths.device)[None, :] >= lengths[:, None] - - +# =============================== +# MAIN +# =============================== def main(): - # =============================== - # 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)) - # =============================== - # CONFIG - # =============================== - base_path = "asl_kaggle" # ← CHANGE THIS TO YOUR ACTUAL PATH + base_path = "asl_kaggle" max_frames = 100 - MIN_SAMPLES_PER_CLASS = 6 # ← important! prevents stratified split crash + MIN_SAMPLES_PER_CLASS = 6 - # =============================== - # DATA LOADING & PROCESSING - # =============================== - print("Loading metadata...") + # --- LOAD DATA --- 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 - ), + pool.imap(partial(process_row, base_path=base_path, max_frames=max_frames), rows), total=len(rows), - desc="Extracting landmarks" + desc="Processing" )) - X_list, y_list = [], [] - for feat, sign in results: + X_list, mask_list, y_list = [], [], [] + for feat, mask, sign in results: if feat is not None: X_list.append(feat) + mask_list.append(mask) y_list.append(sign) if not X_list: - print("No valid sequences found. Check parquet files / paths.") + print("No valid sequences found.") return X = np.stack(X_list) - print(f"Loaded {len(X)} valid sequences | shape: {X.shape}") + masks = np.stack(mask_list) + print(f"Loaded {len(X)} sequences | shape: {X.shape}") - # Global normalization (very important for stability) - print("Before global norm → mean:", X.mean(), "std:", X.std()) - X = np.clip(X, -5.0, 5.0) - 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()) + # --- NORMALIZE only valid frames --- + for i in range(X.shape[0]): + valid_idx = masks[i] + X[i, valid_idx] = (X[i, valid_idx] - X[i, valid_idx].mean(0)) / (X[i, valid_idx].std(0) + 1e-8) - # =============================== - # LABELS - # =============================== + # --- LABELS --- le = LabelEncoder() y = le.fit_transform(y_list) - # Remove classes with too few samples (prevents stratify error) + # Remove rare classes counts = Counter(y) valid_classes = [cls for cls, cnt in counts.items() if cnt >= MIN_SAMPLES_PER_CLASS] - - mask = np.isin(y, valid_classes) - X = X[mask] - y = y[mask] - - # Re-encode labels consecutively (0,1,2,... no gaps) + mask_keep = np.isin(y, valid_classes) + X, masks, y = X[mask_keep], masks[mask_keep], y[mask_keep] le = LabelEncoder() y = le.fit_transform(y) + print(f"{len(X)} samples remain | {len(le.classes_)} classes") - print(f"After filtering: {len(X)} samples remain | {len(le.classes_)} classes") - - # =============================== - # SPLIT - # =============================== - X_train, X_test, y_train, y_test = train_test_split( - X, y, - test_size=0.15, - stratify=y, # should be safe now - random_state=42 + # --- SPLIT --- + X_train, X_test, masks_train, masks_test, y_train, y_test = train_test_split( + X, masks, y, test_size=0.15, stratify=y, random_state=42 ) - # =============================== - # DATASET & LOADERS - # =============================== + # --- DATASETS --- class ASLSequenceDataset(Dataset): - def __init__(self, X, y): + def __init__(self, X, masks, y): self.X = torch.from_numpy(X).float() + self.masks = torch.from_numpy(masks) 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] + return self.X[idx], self.masks[idx], self.y[idx] - train_loader = DataLoader( - ASLSequenceDataset(X_train, y_train), - batch_size=64, - shuffle=True, - num_workers=4, - pin_memory=True - ) + train_loader = DataLoader(ASLSequenceDataset(X_train, masks_train, y_train), + batch_size=64, shuffle=True, num_workers=4, pin_memory=True) + test_loader = DataLoader(ASLSequenceDataset(X_test, masks_test, y_test), + batch_size=96, shuffle=False, 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=X.shape[2], num_classes=len(le.classes_)).to(device) + print(f"Model params: {sum(p.numel() for p in model.parameters()):,}") - # =============================== - # MODEL - # =============================== - model = TransformerASL( - input_dim=63, - num_classes=len(le.classes_), - d_model=192, - nhead=6, - num_layers=4 - ).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=5e-4, weight_decay=1e-4) + optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10) - # =============================== - # TRAIN / EVAL - # =============================== + # --- TRAINING --- + best_acc = 0.0 + patience = 15 + wait = 0 + epochs = 70 + def train_epoch(): model.train() total_loss = 0 correct = total = 0 - - for x, y in tqdm(train_loader, desc="Train"): - x, y = x.to(device), y.to(device) - lengths = (x.abs().sum(dim=2) > 1e-5).sum(dim=1) - mask = create_padding_mask(lengths, x.size(1)) + for x, m, yb in tqdm(train_loader, desc="Train"): + x, m, yb = x.to(device), m.to(device), yb.to(device) + mask = create_padding_mask(m) optimizer.zero_grad(set_to_none=True) logits = model(x, key_padding_mask=mask) - - loss = criterion(logits, y) + loss = criterion(logits, yb) loss.backward() - - grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.8) - + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.8) optimizer.step() total_loss += loss.item() - correct += (logits.argmax(-1) == y).sum().item() - total += y.size(0) - + correct += (logits.argmax(-1) == yb).sum().item() + total += yb.size(0) return total_loss / len(train_loader), correct / total * 100 @torch.no_grad() def evaluate(): model.eval() correct = total = 0 - for x, y in test_loader: - x, y = x.to(device), y.to(device) - lengths = (x.abs().sum(dim=2) > 1e-5).sum(dim=1) - mask = create_padding_mask(lengths, x.size(1)) - + for x, m, yb in test_loader: + x, m, yb = x.to(device), m.to(device), yb.to(device) + mask = create_padding_mask(m) logits = model(x, key_padding_mask=mask) - correct += (logits.argmax(-1) == y).sum().item() - total += y.size(0) + correct += (logits.argmax(-1) == yb).sum().item() + total += yb.size(0) return correct / total * 100 if total > 0 else 0.0 - # =============================== - # TRAINING LOOP - # =============================== - best_acc = 0.0 - patience = 15 - wait = 0 - epochs = 70 - 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+1}/{epochs}] loss: {loss:.4f} | train: {train_acc:.2f}% | test: {test_acc:.2f}%") scheduler.step() - if test_acc > best_acc: best_acc = test_acc wait = 0 @@ -362,8 +295,7 @@ def main(): print("Early stopping") break - print(f"\nBest test accuracy reached: {best_acc:.2f}%") - + print(f"\nBest test accuracy: {best_acc:.2f}%") if __name__ == '__main__': - main() \ No newline at end of file + main()