Adding proper profiler

.github/workflows/pr-ci.yml

@@ -11,7 +11,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  test-model-training:
+  test-model-training-and-prediction-generation:
     # Only run on:
     # - push
     # - PRs where the branch lives in this repo (not a fork)
@@ -23,7 +23,7 @@ jobs:
     # Self hosted runner
     runs-on: [self-hosted]
 
-    # The runner is self hosted on my machine, so this data path will be guarenteed to have this data path
+    # The runner is self hosted on my machine, so this data path will be guaranteed to have this data path
     env:
       CI_DATA_ROOT: /home/sam/repos/hobby-repos/nfl/data/parquet
 
@@ -42,11 +42,11 @@ jobs:
       - name: Test feature creation
         run: uv run src/create_features.py --profile
 
-        # This will output model_params.json
+        # This will output model_params.json (--ci only runs 5 epochs for a single trial)
       - name: Test ray hyperparameter tuning
         run: uv run src/train_transformer.py --tune --ci
 
-        # This will use model_params.json to create best_model.pth
+        # 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
 

.gitignore

@@ -189,8 +189,8 @@ output/**/*.txt
 mlruns/
 mlflow.db
 
-# Ray
-ray_results/
+# Ray and profiler
+log
 
 # Artifacts
 *.pkl

data/.gitignore

@@ -1,2 +1,3 @@
 raw/
-processed/
+processed/
+inference/

pyproject.toml

@@ -59,17 +59,18 @@ where = ["src"]
 [tool.ruff]
 target-version = "py312"
 line-length = 200
+exclude = ["*.ipynb"]
+
+[tool.ruff.lint]
 select = [
   "E",   # general pep8 error rules (formatting)
-  "F",   # pyflakes (logica/correctness)
+  "F",   # pyflakes (logic/correctness)
   "B",   # bugbear (logic traps)
   "I",   # isort-like import rules (import org)
   "UP",  # pyupgrade (modern syntax)
   "ANN", # type annotation rules
 ]
-exclude = ["*.ipynb"]
 
-[tool.ruff.lint]
 ignore = [
   "E203", # makes it compat with Black-style slicing
   "E501", # ignore line length

src/generate_predictions.py

@@ -12,7 +12,7 @@
 import pandas as pd
 import polars as pl
 import torch
-from torch.profiler import ProfilerActivity
+from torch.profiler import ProfilerActivity, profile, record_function
 
 from clean_data import calculate_velocity_components, label_offense_defense_manzone, make_plays_left_to_right, pass_attempt_merging, rotate_direction_and_orientation
 from common.args import parse_args
@@ -21,6 +21,8 @@
 from load_data import RawDataLoader
 from models.transformer import create_transformer_model
 
+# ... at top of file you already import ProfilerActivity, so just add record_function, schedule if needed.
+
 
 def process_week_data_preds(week_number: int, plays: pd.DataFrame) -> pd.DataFrame:
 	"""
@@ -102,6 +104,62 @@ def prepare_tensor(play: pd.DataFrame) -> torch.Tensor:
 	return all_frames_tensor  # Shape: [num_frames, num_players, num_features]
 
 
+def profile_inference(model: torch.nn.Module, device: torch.device) -> None:
+	"""
+	Profiles inference of a single frame using a saved example tensor.
+
+	Inputs:
+	- model: The transformer model to profile.
+	- device: The device to run the model on.
+
+	Outputs:
+	- Saved trace_inference.json file in log/profiler/ to visualize in Chrome/Perfetto.
+	"""
+
+	example_path = PROJECT_ROOT / "data" / "inference" / "example_frame_tensor.pt"
+	if example_path.exists():
+		frame_tensor = torch.load(example_path).to(device)
+
+		# Warmup to trigger torch.compile JIT and CUDA autotuning outside profiler
+		with torch.no_grad():
+			for _ in range(5):
+				_ = model(frame_tensor)
+
+		log_dir = PROJECT_ROOT / "log" / "profiler"
+		log_dir.mkdir(parents=True, exist_ok=True)
+		trace_path = log_dir / "trace_inference.json"
+
+		# Setup profiler
+		activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
+		with (
+			torch.no_grad(),
+			profile(
+				activities=activities,
+				record_shapes=True,
+				profile_memory=True,
+				with_stack=True,
+			) as prof,
+		):
+			# Run several inferences to get stable averages
+			for _ in range(20):
+				with record_function("inference_step"):
+					_ = model(frame_tensor)
+
+		print("Top CPU ops (inference):")
+		print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=30))
+		print("Top CUDA ops (inference):")
+		print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=30))
+
+		# Save Chrome/Perfetto trace
+		prof.export_chrome_trace(str(trace_path))
+		print(f"Inference trace saved to: {trace_path}")
+
+		# Early return so we don't also run the full weekly loop when profiling
+		return
+	else:
+		logging.warning("args.profile is set, but example_frame_tensor.pt not found. Run once without profiling to generate it.")
+
+
 @time_fcn
 def main() -> None:
 	"""
@@ -235,6 +293,10 @@ def main() -> None:
 
 		gc.collect()
 
+	# If profiling flag is set, run profiler
+	if args.profile:
+		profile_inference(model, device)
+
 
 if __name__ == "__main__":
 	main()

src/train_transformer.py

@@ -19,6 +19,7 @@
 from ray.tune import FailureConfig, RunConfig, TuneConfig, Tuner
 from ray.tune.schedulers import ASHAScheduler
 from torch.amp import autocast
+from torch.profiler import ProfilerActivity, profile, record_function, schedule
 from torch.utils.data import DataLoader, TensorDataset
 
 from common.args import parse_args
@@ -62,8 +63,23 @@ def set_seed(seed: int = 42) -> torch.Generator:
 	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:
+	# 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):
+		outputs = model(features)
+		loss = loss_fn(outputs, targets)
+
+	# Backpropagate and optimize
+	# NOTE: We don't want the gradients in less than FP32, PyTorch would handle this for us but being explicit
+	loss.backward()
+	optimizer.step()
+
+	return loss
+
+
 @time_fcn
-def train_epoch(train_loader: DataLoader, model: nn.Module, optimizer: torch.optim.Optimizer, loss_fn: nn.Module, device: torch.device, amp_dtype: torch.dtype) -> float:
+def train_epoch(train_loader: DataLoader, model: nn.Module, optimizer: torch.optim.Optimizer, loss_fn: nn.Module, device: torch.device, amp_dtype: torch.dtype, prof: profile) -> float:
 	"""
 	Runs one training epoch over the provided dataloader.
 
@@ -90,16 +106,13 @@ def train_epoch(train_loader: DataLoader, model: nn.Module, optimizer: torch.opt
 		# Zero gradients
 		optimizer.zero_grad(set_to_none=True)
 
-		# 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):
-			outputs = model(features)
-			loss = loss_fn(outputs, targets)
-
-		# Backpropagate and optimize
-		# NOTE: We don't want the gradients in less than FP32, PyTorch would handle this for us but being explicit
-		loss.backward()
-		optimizer.step()
+		# Profiled training step
+		if prof:
+			with record_function("train_step"):
+				loss = forward_backpropagate_optimize_step(amp_dtype, model, features, targets, loss_fn, optimizer)
+			prof.step()
+		else:
+			loss = forward_backpropagate_optimize_step(amp_dtype, model, features, targets, loss_fn, optimizer)
 
 		# Record running loss
 		running_loss += loss.detach().item() * features.size(0)
@@ -237,6 +250,28 @@ def run_trial(config: dict[str, float | int], args: argparse.Namespace) -> None:
 		pin_memory=True,  # Batches are allocated on page-locked ("pinned") memory on the host, allows GPU driver to perform faster async DMA
 	)
 
+	######################################################################
+	# Create profiler
+	######################################################################
+	if args.profile:
+		log_path = PROJECT_ROOT / "log" / "profiler"
+		log_path.mkdir(parents=True, exist_ok=True)
+		prof = profile(
+			activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+			schedule=schedule(
+				wait=1,  # ignore first step (cold start)
+				warmup=1,  # record but don't save
+				active=3,  # record + save 3 steps
+				repeat=2,  # repeat that pattern
+			),
+			record_shapes=True,
+			profile_memory=True,
+			with_stack=True,
+		)
+		prof.start()
+	else:
+		prof = None
+
 	######################################################################
 	# Train model (and evaluate)
 	######################################################################
@@ -252,9 +287,10 @@ def run_trial(config: dict[str, float | int], args: argparse.Namespace) -> None:
 	if args.ci:
 		config["epochs"] = 5
 
+	# Training loop
 	for epoch in range(int(config["epochs"])):
 		# Train
-		avg_train_loss = train_epoch(train_loader, model, optimizer, loss_fn, device, amp_dtype)
+		avg_train_loss = train_epoch(train_loader, model, optimizer, loss_fn, device, amp_dtype, prof)
 		train_losses.append(avg_train_loss)
 
 		# Validate
@@ -293,6 +329,15 @@ def run_trial(config: dict[str, float | int], args: argparse.Namespace) -> None:
 					logging.info("Early stopping triggered")
 					break
 
+	# Stop profiler
+	if prof:
+		prof.stop()
+		print("Top CPU ops:")
+		print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=30))
+		print("Top GPU (CUDA kernel) ops:")
+		print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=30))
+		prof.export_chrome_trace(str(log_path / "trace_training.json"))
+
 
 @time_fcn
 def run_hpo(args: argparse.Namespace) -> None: