diff --git a/.github/workflows/pr-ci.yml b/.github/workflows/pr-ci.yml index f749552..996a664 100644 --- a/.github/workflows/pr-ci.yml +++ b/.github/workflows/pr-ci.yml @@ -47,8 +47,28 @@ jobs: run: uv run src/train_transformer.py --tune --ci # This will use model_params.json to create transformer.pt (--ci only runs 5 epochs) - - name: Test transformer training - run: uv run src/train_transformer.py --profile --ci + - name: Check deterministic transformer training + run: | + set -euo pipefail + uv run src/train_transformer.py --profile --ci + cp data/training/training_metrics.json /tmp/training_metrics_run1.json + uv run src/train_transformer.py --profile --ci + python3 - <<'PY' + import json + import math + from pathlib import Path + + first = json.loads(Path("/tmp/training_metrics_run1.json").read_text()) + second = json.loads(Path("data/training/training_metrics.json").read_text()) + tol = 1e-5 + keys = ("train_loss", "val_loss", "val_accuracy", "best_val_loss") + for key in keys: + if not math.isclose(first[key], second[key], rel_tol=0.0, abs_tol=tol): + raise SystemExit( + f"Determinism check failed for {key}: {first[key]} vs {second[key]} (tol={tol})" + ) + print("Determinism check passed.") + PY # This will use transformer.pt for inference, will make predictions for all of week 9 - name: Test prediction generation @@ -56,4 +76,4 @@ jobs: # This will use the predictions saved as parquet in stage above, only 1 plot, one animated play, and one confusion matrix is created - name: Test process predictions - run: uv run src/process_predictions.py --profile --ci \ No newline at end of file + run: uv run src/process_predictions.py --profile --ci diff --git a/README.md b/README.md index 9e70b24..1283fa6 100644 --- a/README.md +++ b/README.md @@ -8,114 +8,6 @@ Two models are implemented: Both models yield about 88% accuracy in predicting man or pass coverage BEFORE the snap for all plays in Week 9 of the 2022 NFL season. -# Data + modeling pipeline (what each script does) - -This section maps the end-to-end workflow to the code paths you’ll run. It focuses on the transformer workflow, which is what the repo uses for Week 9 evaluation. - -## 1) Data engineering + feature creation (`src/create_features.py`) - -`create_features.py` orchestrates the full data engineering pipeline. At a high level it: - -1. **Loads raw parquet data** via `RawDataLoader` (`src/load_data.py`): - - Combines 2021 + 2022 `games`, `plays`, and `players` into unified DataFrames. - - Loads weekly tracking files for each season/week requested. - - Normalizes older 2021 tracking data by copying `team` into `club` (2021 lacks `club`). - -2. **Cleans and normalizes tracking data** via `clean_df` in `create_features.py` (using functions in `src/clean_data.py`): - - `rotate_direction_and_orientation`: normalizes direction/orientation so 0° points left-to-right. - - `make_plays_left_to_right`: flips every play so offense always moves left-to-right; creates `x_clean`, `y_clean`, `o_clean`, `dir_clean`, etc. - - `calculate_velocity_components`: derives `v_x`, `v_y` from cleaned speed + direction. - - `pass_attempt_merging`: joins play metadata to flag pass attempts. - - `label_offense_defense_manzone`: joins PFF man/zone labels and identifies defenders (`defense` flag). - - Filters to defenders + pass attempts only, adds `week`, `uniqueId`, `frameUniqueId`, and computes `frames_from_snap`. - - Keeps frames from 15s before snap through 5s after snap (`frames_from_snap` between -150 and 50). - -3. **Consolidates + subsamples frames** via `_process_df` in `create_features.py`: - - Keeps **even frames only** (`frameId % 2 == 0`) for a lighter, consistent sampling rate. - - Selects snap-centered frames for augmentation using `select_augmented_frames`. - -4. **Synthetic data generation (augmentation)** via `data_augmentation`: - - Mirrors selected frames left-to-right across the field (`y_clean` flip + direction flip). - - Appends `_aug` to `frameUniqueId` to keep augmented frames distinct. - -5. **Feature tensorization** via `prepare_frame_data`: - - For each `frameUniqueId`, stacks per-player features into a fixed tensor. - - Features used: `["x_clean", "y_clean", "v_x", "v_y", "defense"]`. - - Target label: `pff_manZone` (mapped Zone→0, Man→1). - -6. **Writes weekly + pooled tensors**: - - Saves per-week tensors: `data/processed/s{season}_w{week}_features.pt` and `..._targets.pt`. - - Concatenates weeks into: - - `features_training.pt` / `targets_training.pt` (2021 W1–8 + 2022 W1–8) - - `features_val.pt` / `targets_val.pt` (2022 W9) - -## 2) Model training (`src/train_transformer.py` + `src/models/transformer.py`) - -`train_transformer.py` loads the prepared tensors and trains the transformer model defined in `src/models/transformer.py`. - -**Model architecture (`src/models/transformer.py`):** -- `ManZoneTransformer` takes `[batch, num_players, 5]` inputs: - - 5 features = `x_clean`, `y_clean`, `v_x`, `v_y`, `defense`. -- Layers: - - BatchNorm over feature dimension. - - Linear + ReLU + LayerNorm embedding to `model_dim`. - - Transformer encoder stack (`num_layers`, `num_heads`, `dim_feedforward`, `dropout`). - - Adaptive average pooling across players (reduces to a per-frame embedding). - - MLP decoder to 2 logits (Zone vs Man). - -**Training flow (`src/train_transformer.py`):** -1. Loads `features_training.pt` / `targets_training.pt` and validation tensors. -2. Builds `DataLoader`s with fixed seeds for reproducible shuffling. -3. Uses `AdamW` optimizer + `CrossEntropyLoss`. -4. Runs mixed precision forward passes (`torch.autocast`) on GPU for speed. -5. Tracks training + validation loss and accuracy each epoch. -6. Saves the best model to `data/training/transformer.pt` (by lowest validation loss). -7. Optional Ray Tune (`--tune`) runs hyperparameter search and writes best config to `data/training/model_params.json`. - -## 3) Prediction generation (`src/generate_predictions.py`) - -`generate_predictions.py` loads the trained transformer and streams Week 9 frames through it: - -1. **Re-loads and cleans tracking data** (same cleaning logic as training): - - `rotate_direction_and_orientation` - - `make_plays_left_to_right` - - `calculate_velocity_components` - - `pass_attempt_merging` - - `label_offense_defense_manzone` - - Adds `uniqueId`, `frameUniqueId`, `frames_from_snap`. - - Filters to pass attempts + defenders, and keeps frames from -150 to +30. - -2. **Builds frame tensors on the fly**: - - For each `frameUniqueId`, builds `[num_frames, num_players, 5]` tensor. - - Runs inference with the compiled transformer. - -3. **Writes prediction outputs**: - - Per-frame predictions saved as `data/inference/week{week}_preds.csv` with: - - `zone_prob`, `man_prob`, `pred`, `actual`. - - Joins predictions back to tracking data and saves: - - `data/inference/tracking_s{season}_w{week}_preds.parquet`. - -## 4) Post-processing + evaluation (`src/process_predictions.py`) - -`process_predictions.py` turns the predictions into evaluation artifacts: - -1. **Accuracy vs. time-from-snap plot**: - - Aggregates frame-level correctness by `seconds_from_snap`. - - Saves `artifacts/plot_s{season}_w{week}_accuracy.png`. - -2. **Play-level animation artifacts**: - - Finds plays with the largest increase in man-coverage probability pre-snap. - - Generates animated GIFs of man-probability vs. time for top plays. - -3. **Confusion matrix**: - - Builds a normalized confusion matrix (Zone vs. Man). - - Saves `artifacts/confusion_matrix.png`. - -4. **ROC curve (optional extension)**: - - `generate_predictions.py` already writes probabilities (`man_prob`, `zone_prob`) per frame. - - If you want ROC/AUC, those probabilities are ready to feed into `sklearn.metrics.roc_curve` and `auc`. - - (Currently, the script outputs confusion matrix + accuracy plot; ROC is straightforward to add on top of the saved probabilities.) - # Demo @@ -466,6 +358,114 @@ Self CUDA time total: 1.592ms An example of the timeline view (with the `trace_inference.json` loaded into perfetto) is shown below: ![profile_inference_timeline](./docs/images/profile_inference_timeline.png) +# Code Overview + +This section maps the end-to-end workflow to the code paths you’ll run. It focuses on the transformer workflow, which is what the repo uses for Week 9 evaluation. + +## 1) Data engineering + feature creation (`src/create_features.py`) + +`create_features.py` orchestrates the full data engineering pipeline. At a high level it: + +1. **Loads raw parquet data** via `RawDataLoader` (`src/load_data.py`): + - Combines 2021 + 2022 `games`, `plays`, and `players` into unified DataFrames. + - Loads weekly tracking files for each season/week requested. + - Normalizes older 2021 tracking data by copying `team` into `club` (2021 lacks `club`). + +2. **Cleans and normalizes tracking data** via `clean_df` in `create_features.py` (using functions in `src/clean_data.py`): + - `rotate_direction_and_orientation`: normalizes direction/orientation so 0° points left-to-right. + - `make_plays_left_to_right`: flips every play so offense always moves left-to-right; creates `x_clean`, `y_clean`, `o_clean`, `dir_clean`, etc. + - `calculate_velocity_components`: derives `v_x`, `v_y` from cleaned speed + direction. + - `pass_attempt_merging`: joins play metadata to flag pass attempts. + - `label_offense_defense_manzone`: joins PFF man/zone labels and identifies defenders (`defense` flag). + - Filters to defenders + pass attempts only, adds `week`, `uniqueId`, `frameUniqueId`, and computes `frames_from_snap`. + - Keeps frames from 15s before snap through 5s after snap (`frames_from_snap` between -150 and 50). + +3. **Consolidates + subsamples frames** via `_process_df` in `create_features.py`: + - Keeps **even frames only** (`frameId % 2 == 0`) for a lighter, consistent sampling rate. + - Selects snap-centered frames for augmentation using `select_augmented_frames`. + +4. **Synthetic data generation (augmentation)** via `data_augmentation`: + - Mirrors selected frames left-to-right across the field (`y_clean` flip + direction flip). + - Appends `_aug` to `frameUniqueId` to keep augmented frames distinct. + +5. **Feature tensorization** via `prepare_frame_data`: + - For each `frameUniqueId`, stacks per-player features into a fixed tensor. + - Features used: `["x_clean", "y_clean", "v_x", "v_y", "defense"]`. + - Target label: `pff_manZone` (mapped Zone→0, Man→1). + +6. **Writes weekly + pooled tensors**: + - Saves per-week tensors: `data/processed/s{season}_w{week}_features.pt` and `..._targets.pt`. + - Concatenates weeks into: + - `features_training.pt` / `targets_training.pt` (2021 W1–8 + 2022 W1–8) + - `features_val.pt` / `targets_val.pt` (2022 W9) + +## 2) Model training (`src/train_transformer.py` + `src/models/transformer.py`) + +`train_transformer.py` loads the prepared tensors and trains the transformer model defined in `src/models/transformer.py`. + +**Model architecture (`src/models/transformer.py`):** +- `ManZoneTransformer` takes `[batch, num_players, 5]` inputs: + - 5 features = `x_clean`, `y_clean`, `v_x`, `v_y`, `defense`. +- Layers: + - BatchNorm over feature dimension. + - Linear + ReLU + LayerNorm embedding to `model_dim`. + - Transformer encoder stack (`num_layers`, `num_heads`, `dim_feedforward`, `dropout`). + - Adaptive average pooling across players (reduces to a per-frame embedding). + - MLP decoder to 2 logits (Zone vs Man). + +**Training flow (`src/train_transformer.py`):** +1. Loads `features_training.pt` / `targets_training.pt` and validation tensors. +2. Builds `DataLoader`s with fixed seeds for reproducible shuffling. +3. Uses `AdamW` optimizer + `CrossEntropyLoss`. +4. Runs mixed precision forward passes (`torch.autocast`) on GPU for speed. +5. Tracks training + validation loss and accuracy each epoch. +6. Saves the best model to `data/training/transformer.pt` (by lowest validation loss). +7. Optional Ray Tune (`--tune`) runs hyperparameter search and writes best config to `data/training/model_params.json`. + +## 3) Prediction generation (`src/generate_predictions.py`) + +`generate_predictions.py` loads the trained transformer and streams Week 9 frames through it: + +1. **Re-loads and cleans tracking data** (same cleaning logic as training): + - `rotate_direction_and_orientation` + - `make_plays_left_to_right` + - `calculate_velocity_components` + - `pass_attempt_merging` + - `label_offense_defense_manzone` + - Adds `uniqueId`, `frameUniqueId`, `frames_from_snap`. + - Filters to pass attempts + defenders, and keeps frames from -150 to +30. + +2. **Builds frame tensors on the fly**: + - For each `frameUniqueId`, builds `[num_frames, num_players, 5]` tensor. + - Runs inference with the compiled transformer. + +3. **Writes prediction outputs**: + - Per-frame predictions saved as `data/inference/week{week}_preds.csv` with: + - `zone_prob`, `man_prob`, `pred`, `actual`. + - Joins predictions back to tracking data and saves: + - `data/inference/tracking_s{season}_w{week}_preds.parquet`. + +## 4) Post-processing + evaluation (`src/process_predictions.py`) + +`process_predictions.py` turns the predictions into evaluation artifacts: + +1. **Accuracy vs. time-from-snap plot**: + - Aggregates frame-level correctness by `seconds_from_snap`. + - Saves `artifacts/plot_s{season}_w{week}_accuracy.png`. + +2. **Play-level animation artifacts**: + - Finds plays with the largest increase in man-coverage probability pre-snap. + - Generates animated GIFs of man-probability vs. time for top plays. + +3. **Confusion matrix**: + - Builds a normalized confusion matrix (Zone vs. Man). + - Saves `artifacts/confusion_matrix.png`. + +4. **ROC curve (optional extension)**: + - `generate_predictions.py` already writes probabilities (`man_prob`, `zone_prob`) per frame. + - If you want ROC/AUC, those probabilities are ready to feed into `sklearn.metrics.roc_curve` and `auc`. + - (Currently, the script outputs confusion matrix + accuracy plot; ROC is straightforward to add on top of the saved probabilities.) + # Contributing ## Linting and Formatting diff --git a/src/train_transformer.py b/src/train_transformer.py index 567cc4d..6819731 100644 --- a/src/train_transformer.py +++ b/src/train_transformer.py @@ -2,6 +2,7 @@ """ Trains and tunes a transformer model for man/zone classification on tracking data. """ +from __future__ import annotations import argparse import json @@ -12,7 +13,6 @@ import numpy as np import ray -import torch import torch.nn as nn from ray import tune from ray.air.integrations.mlflow import MLflowLoggerCallback @@ -28,42 +28,50 @@ from models.transformer import create_transformer_model -def set_seed(seed: int = 42) -> torch.Generator: +def configure_determinism(deterministic: bool, seed: int = 42) -> torch.Generator: """ - Sets Python, NumPy, and Torch seeds for reproducible runs. + Configures PyTorch determinism knobs based on training intent. Inputs: + - deterministic: Whether to enforce fully deterministic execution. - seed: Seed value to apply. Outputs: - generator: Torch generator for deterministic DataLoader shuffling. """ - # Set seeds + import torch + + os.environ["PYTHONHASHSEED"] = str(seed) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) g = torch.Generator() # Creates a generator that fixes the shuffle in torch Dataloader g.manual_seed(seed) - os.environ["PYTHONHASHSEED"] = str(seed) - # PyTorch CPU & CUDA - torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) - # Keep cudnn deterministic, but allow normal algorithms - # torch.backends.cudnn.benchmark = False - # torch.backends.cudnn.deterministic = True - - # DO NOT enforce deterministic algorithms globally - # torch.use_deterministic_algorithms(False) + if deterministic: + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = True + torch.use_deterministic_algorithms(True) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + logging.info("Deterministic training enabled (may reduce throughput).") + else: + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True return g def forward_backpropagate_optimize_step(amp_dtype: torch.dtype, model: nn.Module, features: torch.Tensor, targets: torch.Tensor, loss_fn: nn.Module, optimizer: torch.optim.Optimizer) -> torch.Tensor: + import torch + # For forward pass, allow mixed precision for speed in some operations # NOTE: PyTorch keeps some sensitive ops in FP32 for stability with autocast(device_type="cuda", dtype=amp_dtype): @@ -95,6 +103,8 @@ def train_epoch(train_loader: DataLoader, model: nn.Module, optimizer: torch.opt - avg_train_loss: Mean loss across the epoch. """ + import torch + # Training model.train() running_loss = 0.0 @@ -141,6 +151,8 @@ def validate_epoch(val_loader: DataLoader, model: nn.Module, loss_fn: nn.Module, - val_accuracy: Classification accuracy. """ + import torch + # Validation model.eval() val_running_loss = 0.0 @@ -182,8 +194,7 @@ def run_trial(config: dict[str, float | int], args: argparse.Namespace) -> None: - Trains the model and optionally reports metrics to Ray Tune. """ - # Set seeds - g = set_seed(42) + deterministic = args.profile ###################################################################### # Set things that allow for speed-up in training (valid on Ampere+ GPUs) @@ -194,9 +205,8 @@ def run_trial(config: dict[str, float | int], args: argparse.Namespace) -> None: # Use mixed precision with FP16 for speed for forward pass (see train_epoch) amp_dtype = torch.float16 - # Set TensorFloat-32 (TF32) mode for matmul and cudnn (speeds up training on Ampere+ GPUs with minimal impact on accuracy) - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True + # Configure determinism (full determinism can slow training; only enforced when profiling) + g = configure_determinism(deterministic, seed=42) # Set device device = torch.device("cuda") # this is the Ray-assigned GPU (Ray sets CUDA_VISIBLE_DEVICES) @@ -218,7 +228,9 @@ def run_trial(config: dict[str, float | int], args: argparse.Namespace) -> None: model.parameters(), lr=float(config["lr"]), weight_decay=float(config["weight_decay"]), - fused=True, # All operations on single CUDA kernel for speed + # Fused AdamW uses a single CUDA kernel for speed but can introduce non-deterministic reductions. + # Disable it when deterministic mode is requested. + fused=not deterministic, ) loss_fn = nn.CrossEntropyLoss() @@ -321,6 +333,17 @@ def run_trial(config: dict[str, float | int], args: argparse.Namespace) -> None: best_val_loss = avg_val_loss torch.save(model.state_dict(), best_model_path) + if not args.tune and train_losses: + metrics_output = { + "train_loss": float(train_losses[-1]), + "val_loss": float(val_losses[-1]), + "val_accuracy": float(val_accuracies[-1]), + "best_val_loss": float(best_val_loss), + } + metrics_path = PROJECT_ROOT / "data" / "training" / "training_metrics.json" + with open(metrics_path, "w") as json_file: + json.dump(metrics_output, json_file, indent=4) + # Stop profiler if prof: prof.stop()