#!/usr/bin/env python3 from __future__ import annotations import argparse import copy import time from pathlib import Path import matplotlib.pyplot as plt import numpy as np import torch from torch import nn from torch.utils.data import DataLoader, TensorDataset from fall_detection.experiment import prepare_experiment_data, split_summary from fall_detection.features import CORE_FEATURE_INDICES from fall_detection.metrics import ( choose_f1_threshold, classification_metrics, save_evaluation_plots, save_predictions, ) from fall_detection.models import PoseGRU, PoseTCN from fall_detection.utils import set_seed, write_json def make_loader( features: np.ndarray, labels: np.ndarray, indices: np.ndarray, batch_size: int, shuffle: bool, ) -> DataLoader: dataset = TensorDataset( torch.from_numpy(features[indices]).float(), torch.from_numpy(labels[indices]).float(), ) return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle) @torch.no_grad() def predict(model: nn.Module, loader: DataLoader, device: torch.device) -> np.ndarray: model.eval() probabilities = [] for features, _ in loader: logits = model(features.to(device)) probabilities.append(torch.sigmoid(logits).cpu().numpy()) return np.concatenate(probabilities) def main() -> None: parser = argparse.ArgumentParser(description="Train the pose GRU") parser.add_argument("--dataset", default="data/processed/urfd_pose.npz") parser.add_argument("--config", default="configs/default.yaml") parser.add_argument("--output", type=Path, default=Path("artifacts/experiments/urfd")) parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="auto") parser.add_argument("--seed", type=int) parser.add_argument("--architecture", choices=["gru", "tcn"], default="gru") args = parser.parse_args() config, dataset, features, splits = prepare_experiment_data( args.dataset, args.config, args.output, seed=args.seed ) set_seed(config["seed"]) labels = dataset["labels"] feature_indices = np.arange(features.shape[2], dtype=np.int64) if args.architecture == "tcn": feature_indices = CORE_FEATURE_INDICES features = features[:, :, feature_indices] train_mean = features[splits.train].mean(axis=(0, 1), keepdims=True) train_std = features[splits.train].std(axis=(0, 1), keepdims=True) train_std = np.maximum(train_std, 1e-5) scaled = ((features - train_mean) / train_std).astype(np.float32) if args.device == "auto": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") else: device = torch.device(args.device) if device.type == "cuda" and not torch.cuda.is_available(): raise SystemExit("CUDA was requested but torch.cuda.is_available() is false") training = config["training"] batch_size = int(training["batch_size"]) train_loader = make_loader(scaled, labels, splits.train, batch_size, True) val_loader = make_loader(scaled, labels, splits.val, batch_size, False) test_loader = make_loader(scaled, labels, splits.test, batch_size, False) model_config = config["model"] if args.architecture == "gru": model = PoseGRU( input_size=features.shape[2], hidden_size=int(model_config["hidden_size"]), num_layers=int(model_config["num_layers"]), dropout=float(model_config["dropout"]), ).to(device) model_name = "pose_gru" else: model = PoseTCN( input_size=features.shape[2], channels=int(model_config["tcn_channels"]), dropout=float(model_config["dropout"]), ).to(device) model_name = "pose_tcn" train_labels = labels[splits.train] pos_weight = float(np.sum(train_labels == 0) / max(np.sum(train_labels == 1), 1)) criterion = nn.BCEWithLogitsLoss(pos_weight=torch.tensor(pos_weight, device=device)) optimizer = torch.optim.AdamW( model.parameters(), lr=float(training["learning_rate"]), weight_decay=float(training["weight_decay"]), ) best_loss = float("inf") best_state = None wait = 0 history = {"train_loss": [], "val_loss": []} started = time.perf_counter() for epoch in range(1, int(training["epochs"]) + 1): model.train() train_losses = [] for batch_features, batch_labels in train_loader: batch_features = batch_features.to(device) batch_labels = batch_labels.to(device) optimizer.zero_grad(set_to_none=True) loss = criterion(model(batch_features), batch_labels) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=2.0) optimizer.step() train_losses.append(float(loss.item())) model.eval() val_losses = [] with torch.no_grad(): for batch_features, batch_labels in val_loader: batch_features = batch_features.to(device) batch_labels = batch_labels.to(device) val_losses.append(float(criterion(model(batch_features), batch_labels).item())) train_loss = float(np.mean(train_losses)) val_loss = float(np.mean(val_losses)) history["train_loss"].append(train_loss) history["val_loss"].append(val_loss) print(f"Epoch {epoch:02d}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}") if val_loss < best_loss - 1e-4: best_loss = val_loss best_state = copy.deepcopy(model.state_dict()) wait = 0 else: wait += 1 if wait >= int(training["patience"]): print("Early stopping") break training_seconds = time.perf_counter() - started if best_state is None: raise RuntimeError("Training did not produce a checkpoint") model.load_state_dict(best_state) val_probabilities = predict(model, val_loader, device) threshold = choose_f1_threshold(labels[splits.val], val_probabilities) test_probabilities = predict(model, test_loader, device) metrics = classification_metrics(labels[splits.test], test_probabilities, threshold) metrics.update( { "model": model_name, "architecture": args.architecture, "seed": int(config["seed"]), "device": str(device), "device_name": torch.cuda.get_device_name(device) if device.type == "cuda" else "CPU", "torch_version": torch.__version__, "torch_cuda_version": torch.version.cuda, "epochs_completed": len(history["train_loss"]), "best_val_loss": best_loss, "training_seconds": training_seconds, "parameters": sum(parameter.numel() for parameter in model.parameters()), "dataset": str(args.dataset), "split": split_summary(labels, dataset["groups"], splits), "result_scope": "test set only", } ) output_dir = args.output / model_name output_dir.mkdir(parents=True, exist_ok=True) torch.save( { "state_dict": {key: value.cpu() for key, value in best_state.items()}, "architecture": args.architecture, "input_size": int(features.shape[2]), "hidden_size": int(model_config["hidden_size"]), "num_layers": int(model_config["num_layers"]), "tcn_channels": int(model_config["tcn_channels"]), "dropout": float(model_config["dropout"]), "feature_indices": feature_indices, "sequence_length": int(features.shape[1]), "visibility_threshold": float(config["data"]["visibility_threshold"]), "feature_mean": train_mean.squeeze().astype(np.float32), "feature_std": train_std.squeeze().astype(np.float32), "threshold": threshold, }, output_dir / "model.pt", ) write_json(output_dir / "metrics.json", metrics) save_predictions( output_dir / "predictions.csv", labels[splits.test], test_probabilities, dataset["sources"][splits.test], threshold, ) save_evaluation_plots( labels[splits.test], test_probabilities, threshold, output_dir, model_name ) plt.figure(figsize=(6, 4)) plt.plot(history["train_loss"], label="Train") plt.plot(history["val_loss"], label="Validation") plt.xlabel("Epoch") plt.ylabel("BCE loss") plt.title(f"Qua trinh huan luyen {model_name}") plt.legend() plt.tight_layout() plt.savefig(output_dir / "training_curve.png", dpi=180) plt.close() print( f"{model_name}: F1={metrics['f1']:.3f}, recall={metrics['recall']:.3f}, " f"specificity={metrics['specificity']:.3f}, AUC={metrics['roc_auc']:.3f}" ) if __name__ == "__main__": main()