train.py

import torch
import time
import torch.nn as nn
from torch.utils.data import DataLoader
import numpy as np
import sys
import os
from tqdm import tqdm
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from models.mlp import MLP
from models.gru import GRUModel
from models.rnn import RNNModel
from models.lstm import LSTMModel
from models.stgcn import STGCN
from models.tiny_transformer import TinyTransformerModel
from sklearn.metrics import confusion_matrix, accuracy_score

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)

# -------------------------
# MPJPE & MPJVE Metric
# -------------------------
def mpjpe(pred, target):
    return torch.mean(torch.norm(pred - target, dim=-1))

def mpjve(pred, target):
    pred_vel = pred[:, 1:] - pred[:, :-1]
    target_vel = target[:, 1:] - target[:, :-1]
    return torch.mean(torch.norm(pred_vel - target_vel, dim=-1))

model_name = "stgcn"  
batch_size = 32
num_epochs = 1000
lr = 0.0001

# -------------------------
# Training Loop
# -------------------------
def train(model, train_loader, num_epochs=1000, lr=1e-4, model_name="default_model"):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_forecast = nn.MSELoss()
    loss_class = nn.CrossEntropyLoss()
    
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=10, factor=0.5)

    epoch_losses = []
    epoch_accuracies = [] 
    epoch_mpjpes = [] 
    
    start_time = time.time()

    for epoch in range(num_epochs):
        model.train()
        total_loss, total_loss_f, total_loss_c = 0.0, 0.0, 0.0
        correct, total = 0, 0
        all_mpjpe = []
        all_mpjve = []

        y_true_epoch, y_pred_epoch = [], []

        for src, (trg_forecast, trg_class), _, _ in tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}"):
            src = src.to(device)
            trg_forecast = trg_forecast.to(device)
            trg_class = trg_class.to(device)

            if model_name == "stgcn":
                src = src.permute(0, 3, 1, 2)  
                src = src.unsqueeze(-1)    
            else:
                src = src.view(src.size(0), src.size(1), -1)

            trg_forecast = trg_forecast.view(trg_forecast.size(0), trg_forecast.size(1), -1)

            optimizer.zero_grad()
            forecast_out, class_out = model(src)

            loss_f = loss_forecast(forecast_out, trg_forecast)
            loss_c = loss_class(class_out, torch.argmax(trg_class, dim=1))
            loss = loss_f + loss_c 
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * src.size(0)
            total_loss_f += loss_f.item() * src.size(0)
            total_loss_c += loss_c.item() * src.size(0)

            all_mpjpe.append(mpjpe(forecast_out, trg_forecast).item())
            all_mpjve.append(mpjve(forecast_out, trg_forecast).item())

            preds = torch.argmax(class_out, dim=1)
            labels = torch.argmax(trg_class, dim=1)
            correct += (preds == torch.argmax(trg_class, dim=1)).sum().item()
            total += trg_class.size(0)

            y_pred_epoch.extend(preds.cpu().numpy())
            y_true_epoch.extend(labels.cpu().numpy())

        avg_loss = total_loss / total
        avg_mpjpe = sum(all_mpjpe) / len(all_mpjpe)
        acc = correct / total

        epoch_losses.append(avg_loss)
        epoch_accuracies.append(acc)
        epoch_mpjpes.append(avg_mpjpe)

        print(f"[Epoch {epoch+1}] Train → Loss: {avg_loss:.4f} | Forecast: {(total_loss_f/total):.4f} | "
              f"Class: {(total_loss_c/total):.4f} | MPJPE: {avg_mpjpe:.4f} | Acc: {acc:.4f}")
        
        scheduler.step(avg_loss)

    end_time = time.time()
    total_minutes = (end_time - start_time) / 60
    print(f"[INFO] Total Training Time: {total_minutes:.2f} minutes")

    os.makedirs("results/saved_models", exist_ok=True)
    model_path = f"results/saved_models/{model_name}_{num_epochs}.pt"
    torch.save(model.state_dict(), model_path)
    print(f"[INFO] Model saved to {model_path}")

    os.makedirs("results/plots", exist_ok=True)

    # 1. Loss Plot
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, num_epochs + 1), epoch_losses, color='blue', linewidth=2)
    plt.title("Training Loss (Balanced)", fontsize=14)
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.grid(True)
    plt.savefig(f"results/plots/{model_name}_loss.png")
    plt.close()

    # 2. Accuracy Plot
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, num_epochs + 1), epoch_accuracies, color='green', linewidth=2)
    plt.title("Training Accuracy", fontsize=14)
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.grid(True)
    plt.savefig(f"results/plots/{model_name}_accuracy.png")
    plt.close()

    # 3. MPJPE Plot
    plt.figure(figsize=(10, 6))
    plt.plot(range(1, num_epochs + 1), epoch_mpjpes, color='red', linewidth=2)
    plt.title("Training MPJPE (Forecast Error)", fontsize=14)
    plt.xlabel("Epoch")
    plt.ylabel("MPJPE")
    plt.grid(True)
    plt.savefig(f"results/plots/{model_name}_mpjpe.png")
    plt.close()
    
    print(f"[INFO] Plots saved to results/plots/")

# -------------------------
# Run Training
# -------------------------
if __name__ == "__main__":
    import os
    import pickle
    from utils.dataset import PoseDataset
    from torch.utils.data import DataLoader
    from collections import defaultdict, Counter

    with open(os.path.join(os.path.dirname(__file__), "dataset.pkl"), "rb") as f:
        raw_data = pickle.load(f)

    subjects = raw_data["subjects"]
    motion_types = raw_data["motion_types"]
    unique_subjects = sorted(set(subjects))
    train_subjects = unique_subjects[0:4]

    # Prepare train_data dict
    train_data = {
        "src": [], "trg_forecast": [], "trg_class": [],
        "subjects": [], "motion_types": []
    }
    
    for i, subj in enumerate(subjects):
        if subj in train_subjects:
            train_data["src"].append(raw_data["src"][i])
            train_data["trg_forecast"].append(raw_data["trg_forecast"][i])
            train_data["trg_class"].append(raw_data["trg_class"][i])
            train_data["subjects"].append(subj)
            train_data["motion_types"].append(raw_data["motion_types"][i])

    train_dataset = PoseDataset(train_data, return_subject=True)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

    subject_to_motions = defaultdict(set)

    subject_motion_windows = defaultdict(Counter)

    for src, (trg_forecast, trg_class), subjects, motions in train_loader:
        for subj, motion in zip(subjects, motions):
            subject_motion_windows[subj][motion] += 1

    print("\n[INFO] Sliding window count per subject and motion type:")
    for subj in sorted(subject_motion_windows.keys()):
        print(f"{subj}:")
        for motion, count in subject_motion_windows[subj].items():
            print(f"  - {motion}: {count} sliding windows")

    input_window = len(train_data["src"][0])  # e.g. 15
    keypoints_dim = len(train_data["src"][0][0]) * len(train_data["src"][0][0][0])  # 17 * 2 = 34
    output_window = len(train_data["trg_forecast"][0])
    num_keypoints = len(train_data["src"][0][0])
    num_coords = len(train_data["src"][0][0][0])
    num_forecast_coords = len(train_data["trg_forecast"][0][0][0])

    if model_name == "mlp":
        from models.mlp import MLP
        model = MLP(
            input_size=input_window * keypoints_dim,
            hidden_size=128,
            forecast_window=input_window,
            output_class_size=2
        )
    
    elif model_name == "rnn":
        from models.rnn import RNNModel
        model = RNNModel(
            input_size=34,
            hidden_size=128,
            forecast_window=input_window,
            output_class_size=2
        )

    elif model_name == "lstm":
        from models.lstm import LSTMModel
        model = LSTMModel(
            input_size=34,
            hidden_size=128,
            forecast_window=input_window,
            output_class_size=2
        )

    elif model_name == 'gru':
        from models.gru import GRUModel
        model = GRUModel(
            input_size=34,
            hidden_size=128,
            forecast_window=input_window,
            output_class_size=2
        )

    elif model_name == "tiny_transformer":
        from models.tiny_transformer import TinyTransformerModel
        model = TinyTransformerModel(
            input_size=34,
            forecast_window=input_window,
            output_class_size=2,
            d_model=64,        
            nhead=4,          
            num_layers=2,      
            dim_feedforward=128,
            dropout=0.1
        )

    elif model_name == "stgcn":
        from models.stgcn import STGCN
        graph_args = {'layout': 'coco', 'strategy': 'spatial'}
        model = STGCN(
            in_channels=num_coords,             # Input channels = 4
            num_class=2,
            graph_args=graph_args,
            forecast_window=output_window,
            forecast_channels=num_forecast_coords, # Output channels = 2
            edge_importance_weighting=True
        )

    else:
        raise ValueError(f"Unknown model name: {model_name}")

    model.to(device)

    print(f"Using model: {model_name.upper()}")
    train(model, train_loader, num_epochs=num_epochs, lr=lr, model_name=model_name)