train_fairseq.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import pandas as pd
import numpy as np
import random
from torch.utils.data import Dataset, DataLoader
import pytorch_lightning as pl
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from fairseq_ode_transformer import LandmarkODETransformerEncoder
import argparse
from types import SimpleNamespace
from fairseq.data.dictionary import Dictionary


torch.set_grad_enabled(True)
# Try changing the tensor format
torch.set_default_dtype(torch.float32)
# Also try this flag
torch.backends.cuda.matmul.allow_tf32 = False
# -------------------
# Seed Setting Function
# -------------------
def set_seed(seed):
    """Sets the seed for all random generators"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    # Deterministic operasyonları etkinleştir
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ['PYTHONHASHSEED'] = str(seed)
    print(f"Random seed {seed} olarak ayarlandı.")

# -------------------
# Argument Parsing
# -------------------
def parse_args():
    parser = argparse.ArgumentParser(description='Sign Language Transformer Training')
    parser.add_argument('--num_frames', type=int, default=16, help='Number of frames to sample from each video')
    parser.add_argument('--d_model', type=int, default=512, help='Dimension of the model')
    parser.add_argument('--nhead', type=int, default=4, help='Number of heads in self-attention')
    parser.add_argument('--num_layers', type=int, default=2, help='Number of transformer layers')
    parser.add_argument('--dropout', type=float, default=0.2, help='Dropout rate')
    parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
    parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility')

    # Yeni argümanlar
    parser.add_argument('--encoder_normalize_before', type=lambda x: str(x).lower() in ['1','true','yes'],
                        default=False, help='Encoder pre-layer norm kullanılsın mı (true/false)')
    parser.add_argument('--enc_calculate_num', type=int, default=2, help='enc_calculate_num değeri')
    parser.add_argument('--rk_type', type=str, default='standard',
                        choices=['standard', 'initialization', 'learnable','residual','none'], help='RK integrasyon tipi')
    parser.add_argument('--encoder_history_type', type=str, default='dense',
                        choices=['dense', "none", 'residual'], help='Encoder history tipi')
    return parser.parse_args()

# -------------------
# Dataset
# -------------------
class SignDataset(Dataset):
    def __init__(self, csv_path, data_dir, num_frames=16):
        """
        csv_path: CSV dosyasının yolu (ör. train_labels.csv)
        data_dir: .pt dosyalarının bulunduğu dizin (ör. data/train/)
        num_frames: Her video için örneklenecek kare sayısı
        """
        self.df = pd.read_csv(csv_path, header=None)  # sütun isimleri yok
        self.data_dir = data_dir
        self.num_frames = num_frames  # Artık dışarıdan parametre olarak alınıyor
        self.feature_dim = 300  # 75 x 4
        self.eos_token = torch.ones(1, self.feature_dim)  # EOS token (1'lerden oluşan bir vektör)

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        filename = self.df.iloc[idx, 0]  # 1. sütun = dosya adı
        label = int(self.df.iloc[idx, 1])  # 2. sütun = etiket
        filename = filename + "_color.pt"
        file_path = os.path.join(self.data_dir, filename)
        
        # Video verilerini yükle
        x = torch.load(file_path)  # (frame_count, 75, 4)
        x = x.view(x.shape[0], -1)  # (frame_count, 300)
        
        # Belirtilen kare sayısına örnekle
        total_frames = x.shape[0]
        if total_frames >= self.num_frames:
            # num_frames kare almak için eşit aralıklı indeksler hesapla
            indices = np.linspace(0, total_frames - 1, self.num_frames, dtype=int)
            sampled_frames = x[indices]
        else:
            # Video num_frames kareden kısaysa, mevcut kareleri al
            sampled_frames = x
            # Kalan kareleri sıfır ile doldur
            padding = torch.zeros(self.num_frames - total_frames, self.feature_dim)
            sampled_frames = torch.cat([sampled_frames, padding], dim=0)
        
        # EOS token ekle (son kare olarak)
        x_with_eos = torch.cat([sampled_frames, self.eos_token], dim=0)  # (num_frames+1, 300)
        
        return x_with_eos.float(), torch.tensor(label, dtype=torch.long)

# -------------------
# Transformer Model
# -------------------
class SignLanguageTransformer(pl.LightningModule):
    def __init__(self, num_classes, num_frames=16, d_model=512, nhead=4, num_layers=2,
                 dropout=0.2, lr=1e-4,
                 encoder_normalize_before=False, enc_calculate_num=2,
                 rk_type="standard", encoder_history_type="dense"):
        super().__init__()
        self.save_hyperparameters()
        input_dim = 300
        self.embedding = nn.Linear(input_dim, d_model)
        self.positional_encoding = nn.Parameter(torch.randn(num_frames + 1, d_model))

        args = SimpleNamespace(
            encoder_embed_dim=d_model,
            encoder_ffn_embed_dim=d_model * 4,
            encoder_attention_heads=nhead,
            max_source_positions=512,
            encoder_layers=num_layers,
            dropout=dropout,
            relu_dropout=dropout,
            attention_dropout=dropout,
            encoder_normalize_before=encoder_normalize_before,
            max_relative_length=-1,
            no_token_positional_embeddings=True,
            encoder_learned_pos=False,
            enc_calculate_num=enc_calculate_num,
            rk_type=rk_type,
            use_word_dropout=False,
            word_dropout=0.1,
            encoder_history_type=encoder_history_type,
            encoder_integration_type="avg",
            decoder_history_type="dense",
            decoder_normalize_before=False,
            decoder_layers=0
        )
        
        # Dummy sözlük ve embedding - bunlar gerçekten kullanılmayacak
        dictionary = Dictionary()
        dictionary.add_symbol("<pad>")
        if hasattr(dictionary, "pad_to_multiple_"):
            dictionary.pad_to_multiple_(1)
        else:
            dictionary.finalize(padding_factor=1)
        embed_tokens = nn.Embedding(len(dictionary), args.encoder_embed_dim, padding_idx=dictionary.pad())
        
        # Özel transformer sınıfını kullan
        self.transformer = LandmarkODETransformerEncoder(
            args, 
            dictionary, 
            embed_tokens,
            input_dim=d_model  # Embedding boyutu
        )
        
        self.fc = nn.Linear(d_model, num_classes)
        self.lr = lr
    
    def forward(self, x):
        # x boyutu: [batch_size, num_frames+1, 300]
        x_embedded = self.embedding(x)  # [batch_size, num_frames+1, d_model]
        x_embedded = x_embedded + self.positional_encoding  # Pozisyon kodlaması ekle
        
        # Her örnek için uzunluğu hesapla
        batch_size = x.size(0)
        seq_len = x.size(1)
        src_lengths = torch.full((batch_size,), seq_len, device=x.device)
        
        # Doğrudan float özelliklerini geçir
        encoder_out = self.transformer(x_embedded, src_lengths)
        
        # Çıktı işleme
        out = encoder_out['encoder_out']  # T x B x C
        out = out.transpose(0, 1)  # B x T x C
        out = out.mean(dim=1)  # Global average pooling
        
        return self.fc(out)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()
        self.log("train_loss", loss, prog_bar=True)
        self.log("train_acc", acc, prog_bar=True)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()
        self.log("val_loss", loss, prog_bar=True)
        self.log("val_acc", acc, prog_bar=True)

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        scheduler = {
            'scheduler': torch.optim.lr_scheduler.ReduceLROnPlateau(
                optimizer, mode='min', factor=0.5, patience=3
            ),
            'monitor': 'val_loss',  # Validation loss'a göre LR ayarlanacak
            'interval': 'epoch',
            'frequency': 1
        }
        return [optimizer], [scheduler]

# -------------------
# Training Script
# -------------------
if __name__ == "__main__":
    # Parse command line arguments
    args = parse_args()
    
    # Random seed'i ayarla
    set_seed(args.seed)
    
    # Argümanları görüntüle
    print(f"Training with parameters: num_frames={args.num_frames}, d_model={args.d_model}, "
          f"nhead={args.nhead}, num_layers={args.num_layers}, dropout={args.dropout}, lr={args.lr}, seed={args.seed}, "
          f"encoder_normalize_before={args.encoder_normalize_before}, enc_calculate_num={args.enc_calculate_num}, "
          f"rk_type={args.rk_type}, encoder_history_type={args.encoder_history_type}")
    
    # sample paths
    train_csv = "train_labels.csv"
    val_csv   = "val_labels.csv"
    train_dir = "/home/omer/Masaüstü/datasets/AUTSL_medipipe_landmarks/train"
    val_dir   = "/home/omer/Masaüstü/datasets/AUTSL_medipipe_landmarks/validation"

    # Dataset oluşturulurken num_frames parametresini ilet
    train_dataset = SignDataset(train_csv, train_dir, num_frames=args.num_frames)
    val_dataset   = SignDataset(val_csv, val_dir, num_frames=args.num_frames)

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
    val_loader   = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4)

    # Modeli komut satırı argümanlarıyla oluştur
    if args.encoder_history_type == "none":
        history_type = None
    else:
        history_type = args.encoder_history_type
    model = SignLanguageTransformer(
        num_classes=226,
        num_frames=args.num_frames,
        d_model=args.d_model,
        nhead=args.nhead,
        num_layers=args.num_layers,
        dropout=args.dropout,
        lr=args.lr,
        encoder_normalize_before=args.encoder_normalize_before,
        enc_calculate_num=args.enc_calculate_num,
        rk_type=args.rk_type,
        encoder_history_type=history_type
    )

    # Checkpoint settings
    checkpoint = ModelCheckpoint(monitor="val_acc", mode="max", save_top_k=1)
    
    # Add early stopping callback
    early_stopping = EarlyStopping(
        monitor='val_loss',
        patience=7,  # 7 epoch boyunca gelişme olmazsa durdur
        mode='min'
    )

    trainer = pl.Trainer(
        max_epochs=100,
        accelerator="gpu" if torch.cuda.is_available() else "cpu",
        devices=1,
        callbacks=[checkpoint, early_stopping],
        deterministic=True  # Deterministic davranışı aktif et
    )

    trainer.fit(model, train_loader, val_loader)