Speeding up training

.github/workflows/pr-ci.yml

@@ -39,12 +39,17 @@ jobs:
       - name: Sync uv
         run: uv sync --locked
 
-      - name: Create features
-        run: uv run src/create_features.py
+      - name: Test feature creation
+        run: uv run src/create_features.py --profile
 
-      - name: Train transformer
-        run: uv run src/train_transformer.py
+        # This will output model_params.json
+      - name: Test ray hyperparameter tuning
+        run: uv run src/train_transformer.py --tune --ci
 
-      - name: Make week 9 predictions
-        run: uv run src/generate_predictions.py
+        # This will use model_params.json to create best_model.pth
+      - name: Test transformer training
+        run: uv run src/train_transformer.py --profile --ci
 
+        # This will use best_model.pth for inference
+      - name: Test prediction generation
+        run: uv run src/generate_predictions.py --profile

.gitignore

@@ -190,7 +190,7 @@ mlruns/
 mlflow.db
 
 # Ray
-log/
+ray_results/
 
 # Artifacts
 *.pkl

README.md

@@ -114,10 +114,16 @@ uv run src/create_features.py
 # Train the model
 uv run src/train_transformer.py
 
+# Alternatively, you can run hyper-parameter optimization which will save the best model and parameters with
+uv run src/train_transformer.py --tune
+
 # Run inference on week 9
 uv run src/generate_predictions.py
 
 # From this point you can explore the notebooks/predictions.ipynb, adjust the data path as neccesary and see the animations and plots
+
+# If you would like to time and of the runs you can append --profile, for example
+uv run src/train_transformer.py --profile
 ```
 
 ## Activating MLflow
@@ -133,4 +139,16 @@ mlflow ui \
   --port 5000
 ```
 
-The last command will start the MLflow GUI at your local host loopback on port 5000: http://127.0.0.1:5000
+The last command will start the MLflow GUI at your local host loopback on port 5000: http://127.0.0.1:5000. Viewing MLflow will look something like this (note our results are after a 500 trial hyper-parameter optimization run with Ray Tune)
+
+![mlflow_tune](./docs/images/mlflow_tune.png)
+
+## Future Work
+
+For model improvements I'm interested in switching to a flash attention layer which can accelerate training
+
+I also want to build up per team models, do ANOVA, then highlight specific players that are giving tells on coverages
+
+Want to try out the model on a non-binary problem and try distinguishing between distinct coverage types
+
+See if I can identify blitzer (is the blitzer player identified in the dataset?)

data/training/model_params.json

@@ -1,9 +1,11 @@
 {
-    "epochs": 50,
-    "model_dim": 64,
+    "model_dim": 96,
     "num_heads": 2,
-    "num_layers": 4,
-    "dropout": 0.1,
-    "batch_size": 64,
-    "checkpoint_freq": 10
-}
+    "num_layers": 2,
+    "dropout": 0.0,
+    "multiplier": 2,
+    "lr": 0.00013702308673732994,
+    "weight_decay": 1.128209328073835e-05,
+    "batch_size": 19200,
+    "epochs": 200
+}

src/common/args.py

@@ -18,4 +18,10 @@ def parse_args():
         action="store_true",
         help="Enable @time_fcn timing decorators for profiling",
     )
+
+    parser.add_argument(
+        "--ci",
+        action="store_true",
+        help="Enable CI mode with reduced epochs for faster training during pipeline",
+    )
     return parser.parse_args()

src/generate_predictions.py

@@ -85,20 +85,19 @@ def process_week_data_preds(week_number, plays):
 
 
 def prepare_tensor(play, num_players=22, num_features=5):
+    features = ['x_clean', 'y_clean', 'v_x', 'v_y', 'defense']
+    play_data = play[features + ['frameId']]
+    play_data = play_data.sort_values(by='frameId')
 
-  features = ['x_clean', 'y_clean', 'v_x', 'v_y', 'defense']
-  play_data = play[features + ['frameId']]
-  play_data = play_data.sort_values(by='frameId')
-
-  frames = (
+    frames = (
     play_data
     .groupby('frameId')[features]
     .apply(lambda g: g.to_numpy())
-  )
-  all_frames_tensor = np.stack(frames.to_list())  # Shape: [num_frames, num_players, num_features]
-  all_frames_tensor = torch.tensor(all_frames_tensor, dtype=torch.float32)
+    )
+    all_frames_tensor = np.stack(frames.to_list())  # Shape: [num_frames, num_players, num_features]
+    all_frames_tensor = torch.tensor(all_frames_tensor, dtype=torch.float32)
 
-  return all_frames_tensor  # Shape: [num_frames, num_players, num_features]
+    return all_frames_tensor  # Shape: [num_frames, num_players, num_features]
 
 
 @time_fcn
@@ -117,35 +116,52 @@ def main():
         set_time_decorators_enabled(False)
         logging.info("Timing decorators disabled")
 
+    # 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
+
     # Get model params and set device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     with open(PROJECT_ROOT / "data" / "training" / "model_params.json", 'r') as file:
-        model_params_map = json.load(file)
+        config = json.load(file)
     model = ManZoneTransformer(
-        feature_len=5,                  # num of input features (x, y, v_x, v_y, defense)
-        model_dim=model_params_map["model_dim"],
-        num_heads=model_params_map["num_heads"],
-        num_layers=model_params_map["num_layers"],
-        dim_feedforward=int(model_params_map["model_dim"]) * int(model_params_map["num_layers"]),
-        dropout=model_params_map["dropout"],
-        output_dim=2      # man or zone classification
-    ).to(device)
+        feature_len=5,                                                          # num of input features (x, y, v_x, v_y, defense)
+        model_dim=int(config["model_dim"]),                                     # from ray tune or loaded
+        num_heads=int(config["num_heads"]),                                     # from ray tune or loaded
+        num_layers=int(config["num_layers"]),                                   # from ray tune or loaded
+        dim_feedforward=int(config["model_dim"]) * int(config["multiplier"]),   # from ray tune or loaded
+        dropout=float(config["dropout"]),                                       # from ray tune or loaded
+        output_dim=2                                                            # man or zone classification
+    )
+    # Move model to device (GPU)
+    model = model.to(device)
+
+    # Compile with fullgraph
+    model = torch.compile(
+        model,
+        mode="default",        # default, reduces overhead, generally stable
+        fullgraph=True         # enable full graph fusion, helps reduce kernel launch overhead
+    )
+
+    # Load model
     model.load_state_dict(torch.load(PROJECT_ROOT / "data" / "training" / "best_model.pth", map_location=device))
+
+    # Set to eval mode
     model.eval()
 
     # Load data
     rawLoader = RawDataLoader()
-    games_df, plays_df, players_df, location_data_df = rawLoader.get_data(weeks=[i for i in range(1, 10)])
+    _, plays_df, _, _ = rawLoader.get_data(weeks=list(range(1, 10)))
 
     # Process + predict one week at a time (keeps RAM low)
     for week_eval in [9]:
         week_df = process_week_data_preds(week_eval, plays_df)
 
-        # filter early to shrink memory
+        # Filtering early to shrink memory
         week_df = week_df[(week_df['club'] != 'football') & (week_df['passAttempt'] == 1)].copy()
 
-        # Polars convert optional; if you like the speed, keep it—otherwise skip to use pandas only
-        tracking_df_polars = pl.DataFrame(week_df)  # or comment this out and use pandas below
+        # Polars convert optional for speed
+        tracking_df_polars = pl.DataFrame(week_df)
 
         # Stream predictions to CSV in batches
         weekly_predictions_path_csv = SAVE_DIR / f"week{week_eval}_preds.csv"