diff --git a/deep_quoridor/coding-agents/onnx_save_optimization_plan.md b/deep_quoridor/coding-agents/onnx_save_optimization_plan.md new file mode 100644 index 00000000..6cd378cd --- /dev/null +++ b/deep_quoridor/coding-agents/onnx_save_optimization_plan.md @@ -0,0 +1,44 @@ +# ONNX Save Optimization Plan: Cache ONNX Graph, Update Weights Only + +## Problem + +`save_model_onnx()` calls `torch.onnx.export()` on every save, re-tracing the full +computation graph each time even though the architecture never changes. + +## Proposed Fix + +Run the full export once, cache the ONNX protobuf in memory, then for all subsequent +saves replace only the weight tensors (ONNX initializers) and re-serialize — skipping +graph tracing entirely. + +## Steps + +- [ ] Write plan to file — commit: `vibe: add onnx save optimization plan` +- [ ] Baseline benchmark — run `test_onnx_export.yaml` with `model_save_timing: true`, record per-save timings +- [ ] Verify `onnx` package is available — check `requirements.txt`, add if missing +- [ ] Implement cached ONNX save in `alphazero.py`: + - In `__init__`: add `self._onnx_proto = None` and `self._onnx_init_name_to_idx: dict = {}` + - First call to `save_model_onnx`: run existing `torch.onnx.export(...)`, then load proto and build name→index map + - Subsequent calls: update initializers directly from `state_dict()`, then `onnx.save()` +- [ ] Commit functional change: `vibe: cache ONNX proto and update initializers on subsequent saves` +- [ ] Post-implementation benchmark — run same test YAML, record per-save timings +- [ ] Write timing results to `onnx_save_optimization_results.md` — commit: `vibe: add onnx save optimization benchmark results` +- [ ] Formatting/linting — run ruff, commit: `vibe: formatting for onnx save optimization` + +## Implementation Notes + +- `do_constant_folding=True` is kept; PyTorch's exporter preserves all learned parameters + as explicit ONNX initializers regardless, so the name mapping is complete for all + trainable weights. +- "Cache proto + update initializers" chosen over `torch.jit.trace` reuse — the latter + still rebuilds the ONNX graph on each export; the former bypasses graph construction + entirely after the first save. +- Existing `model_save_timing` flag used for benchmarking — no new standalone script needed. +- No strict numerical output comparison. +- Two separate commits per change (functional + formatting) as per spec. + +## Verification + +- Confirm `model_save_timing` logs show reduced time on saves 2+ +- Load each saved `.onnx` with `onnx.checker.check_model(path)` to confirm structural validity +- Run one inference through `onnxruntime` per saved model to confirm no crash diff --git a/deep_quoridor/coding-agents/onnx_save_optimization_results.md b/deep_quoridor/coding-agents/onnx_save_optimization_results.md new file mode 100644 index 00000000..4d471a61 --- /dev/null +++ b/deep_quoridor/coding-agents/onnx_save_optimization_results.md @@ -0,0 +1,111 @@ +# ONNX Save Optimization Results + +## Setup + +- Experiment config: `deep_quoridor/experiments/B5W3/test_onnx_export.yaml` +- `finish_after: 10 models`, `model_save_timing: true`, `save_onnx: true` +- Tested on both MLP and ResNet network types +- Timing reported by trainer covers both PyTorch `.pt` save and ONNX save together + +--- + +## MLP Network + +### Baseline (before) + +| Save # | Time (s) | +|--------|----------| +| 1 | 1.0424 | +| 2 | 0.9356 | +| 3 | 1.0927 | +| 4 | 1.0188 | +| 5 | 1.0655 | + +**Average per save: ~1.03 s** · **Total for 10 saves: ~10.3 s** + +### After (cached proto) + +| Save # | Time (s) | +|--------|----------| +| 2 | 0.0052 | +| 3 | 0.0041 | +| 4 | 0.0041 | +| 5 | 0.0041 | +| 6–11 | ~0.0041 | + +**Average per save (saves 2+): ~0.0043 s** · **Total for 10 saves: ~0.043 s** + +### MLP Comparison + +| Metric | Before | After | Speedup | +|------------------------|-----------|-----------|-----------| +| First save | ~1.03 s | ~1.03 s | 1× | +| Subsequent saves (avg) | ~1.03 s | ~0.0043 s | **~240×** | +| Total for 10 saves | ~10.3 s | ~0.043 s | **~240×** | + +--- + +## ResNet Network + +### Baseline (before) + +| Save # | Time (s) | +|--------|----------| +| 1 | 1.0264 | +| 2 | 0.9065 | +| 3 | 1.0306 | +| 4 | 0.9120 | +| 5 | 1.0329 | +| 6 | 0.9685 | +| 7 | 1.0510 | +| 8 | 0.9162 | +| 9 | 1.0497 | +| 10 | 1.0103 | + +**Average per save: ~0.987 s** · **Total for 10 saves: ~9.87 s** + +### After (cached proto) + +| Save # | Time (s) | +|--------|----------| +| 2 | 0.0050 | +| 3 | 0.0043 | +| 4 | 0.0045 | +| 5 | 0.0045 | +| 6 | 0.0111 | +| 7 | 0.0064 | +| 8 | 0.0042 | +| 9 | 0.0040 | +| 10 | 0.0039 | +| 11 | 0.0039 | + +**Average per save (saves 2+): ~0.0052 s** · **Total for 10 saves: ~0.052 s** + +### ResNet Comparison + +| Metric | Before | After | Speedup | +|------------------------|-----------|-----------|-----------| +| First save | ~0.99 s | ~0.99 s | 1× | +| Subsequent saves (avg) | ~0.99 s | ~0.0052 s | **~190×** | +| Total for 10 saves | ~9.87 s | ~0.052 s | **~190×** | + +--- + +## Verification + +- All 11 MLP `.onnx` files passed `onnx.checker.check_model()` ✅ +- All 11 ResNet `.onnx` files passed `onnx.checker.check_model()` ✅ +- All models ran inference via `onnxruntime` without error ✅ +- Output shapes: `policy_logits=(1, 57)`, `value=(1, 1)` — correct for 5×5 Quoridor ✅ +- Value outputs in reasonable range for both architectures ✅ + +## Implementation Summary + +- `_onnx_proto = None` and `_onnx_init_name_to_idx = {}` added to `AlphaZeroAgent.__init__` +- First call to `save_model_onnx`: full `torch.onnx.export()` as before, then loads the + written file with `onnx.load()` and builds a name→index map over all graph initializers + (21 for MLP, 21 for ResNet with `num_blocks=2, num_channels=32`). +- Subsequent calls: iterates `state_dict()`, looks up each weight's index in the map, + calls `CopyFrom(onnx.numpy_helper.from_array(...))` in place, then `onnx.save()`. + Graph tracing is skipped entirely. +- Works identically for both MLP and ResNet architectures. diff --git a/deep_quoridor/src/agents/alphazero/alphazero.py b/deep_quoridor/src/agents/alphazero/alphazero.py index 3b35efa7..b059ab91 100644 --- a/deep_quoridor/src/agents/alphazero/alphazero.py +++ b/deep_quoridor/src/agents/alphazero/alphazero.py @@ -11,6 +11,8 @@ import numpy as np import torch +import wandb +import wandb.wandb_run from agents.alphazero.mcts import MCTS, QuoridorKey from agents.alphazero.nn_evaluator import NNConfig, NNEvaluator from agents.core import TrainableAgent @@ -19,9 +21,6 @@ from utils import Timer, get_initial_random_seed, my_device, resolve_path from utils.subargs import SubargsBase -import wandb -import wandb.wandb_run - @dataclass class AlphaZeroParams(SubargsBase): @@ -233,9 +232,7 @@ def __init__( self.action_encoder = ActionEncoder(board_size) nn_config = NNConfig.from_alphazero_params(params) - self.evaluator = NNEvaluator( - self.action_encoder, self.device, nn_config, params.max_cache_size - ) + self.evaluator = NNEvaluator(self.action_encoder, self.device, nn_config, params.max_cache_size) self._fetch_model_from_wandb_and_update_params() self._resolve_and_load_model() @@ -245,9 +242,7 @@ def __init__( # Estimate the typical number of valid moves and use that to choose the alpha parameter of dirichlet noise. # (see parameters class for more explanation of how to choose alpha) if params.mcts_noise_alpha is None: - max_valid_wall_actions = ( - 0 if max_walls == 0 else self.action_encoder.num_actions - board_size**2 - ) + max_valid_wall_actions = 0 if max_walls == 0 else self.action_encoder.num_actions - board_size**2 typical_num_valid_actions = float(np.mean([3, max_valid_wall_actions])) mcts_noise_alpha = 10.0 / typical_num_valid_actions else: @@ -306,9 +301,10 @@ def __init__( # Any game state whose hash has a least significant byte that is in this set is part of the # test set. The test set is the same for the entire run. - self.test_set_lsbs = set( - random.sample(range(256), int(round(256 * params.test_ratio))) - ) + self.test_set_lsbs = set(random.sample(range(256), int(round(256 * params.test_ratio)))) + + self._onnx_proto = None + self._onnx_init_name_to_idx: dict[str, int] = {} def set_wandb_run(self, wandb_run: wandb.wandb_run.Run): self.wandb_run = wandb_run @@ -333,9 +329,7 @@ def _fetch_model_from_wandb_and_update_params(self): artifact = self.wandb_artifact() if artifact is None: return - local_filename = resolve_path( - self.params.wandb_dir, self.wandb_local_filename(artifact) - ) + local_filename = resolve_path(self.params.wandb_dir, self.wandb_local_filename(artifact)) self.params.model_filename = str(local_filename) @@ -349,13 +343,9 @@ def _fetch_model_from_wandb_and_update_params(self): artifact_dir = artifact.download(root=tmpdir) # NOTE: This picks the first .pt file it finds in the artifact - tmp_filename = next( - Path(artifact_dir).glob(f"**/*.{self.get_model_extension()}"), None - ) + tmp_filename = next(Path(artifact_dir).glob(f"**/*.{self.get_model_extension()}"), None) if tmp_filename is None: - raise FileNotFoundError( - f"No model file found in artifact {artifact.name}" - ) + raise FileNotFoundError(f"No model file found in artifact {artifact.name}") shutil.copyfile(tmp_filename, local_filename) @@ -370,9 +360,7 @@ def _resolve_and_load_model(self): if self.params.training_mode: return - print( - "WARNING: no initial model provided usign a filename or wandb, so using random initialized model" - ) + print("WARNING: no initial model provided using a filename or wandb, so using random initialized model") return # If it's not training mode, we definitely need to load a pretrained model, so try the # default path for local files @@ -435,7 +423,14 @@ def save_model(self, path): print(f"AlphaZero model saved to {path}") def save_model_onnx(self, path): - """Export the model to ONNX format.""" + """Export the model to ONNX format. + + On the first call, performs a full torch.onnx.export() and caches the ONNX + protobuf in memory. On subsequent calls, only updates the weight tensors + (initializers) in the cached proto and re-serializes, skipping graph tracing. + """ + import onnx + import onnx.numpy_helper import torch.onnx # Create directory for saving models if it doesn't exist @@ -444,40 +439,68 @@ def save_model_onnx(self, path): # Set the network to evaluation mode self.evaluator.network.eval() - # Create a dummy input tensor with the correct shape - # The shape depends on the network type - network = self.evaluator.network - if ( - hasattr(network, "__class__") - and network.__class__.__name__ == "ResnetNetwork" - ): - # ResNet expects input of shape (batch_size, 5, input_size, input_size) - # NOTE: input_size is board_size * 2 + 3, which is the dimension of the combined grid input, not the original board size - dummy_input = torch.randn( - 1, 5, network.input_size, network.input_size, device=self.device + if self._onnx_proto is None: + # First save: run full export and build the initializer name→index cache. + network = self.evaluator.network + if hasattr(network, "__class__") and network.__class__.__name__ == "ResnetNetwork": + # ResNet expects input of shape (batch_size, 5, input_size, input_size) + # NOTE: input_size is board_size * 2 + 3, which is the dimension of the combined grid input, not the original board size + dummy_input = torch.randn(1, 5, network.input_size, network.input_size, device=self.device) + else: + # MLP expects input of shape (batch_size, input_size) + dummy_input = torch.randn(1, network.input_size, device=self.device) + + # Export the model with opset 17 (widely supported, avoids version conversion issues) + torch.onnx.export( + self.evaluator.network, + (dummy_input,), # Args must be a tuple + str(path), + export_params=True, + opset_version=17, # Use 17 instead of 11 to avoid conversion issues + do_constant_folding=True, + input_names=["input"], + output_names=["policy_logits", "value"], + dynamic_axes={ + "input": {0: "batch_size"}, + "policy_logits": {0: "batch_size"}, + "value": {0: "batch_size"}, + }, + external_data=False, # Don't use external data format for simplicity; model should be small enough to fit in a single file + ) + + self._onnx_proto = onnx.load(str(path)) + state_dict_names = set(self.evaluator.network.state_dict().keys()) + self._onnx_init_name_to_idx = { + init.name: idx + for idx, init in enumerate(self._onnx_proto.graph.initializer) + if init.name in state_dict_names + } + print( + f"AlphaZero model exported to ONNX at {path} " + f"({len(self._onnx_init_name_to_idx)} initializers cached)" ) else: - # MLP expects input of shape (batch_size, input_size) - dummy_input = torch.randn(1, network.input_size, device=self.device) - - # Export the model with opset 17 (widely supported, avoids version conversion issues) - torch.onnx.export( - self.evaluator.network, - (dummy_input,), # Args must be a tuple - path, - export_params=True, - opset_version=17, # Use 17 instead of 11 to avoid conversion issues - do_constant_folding=True, - input_names=["input"], - output_names=["policy_logits", "value"], - dynamic_axes={ - "input": {0: "batch_size"}, - "policy_logits": {0: "batch_size"}, - "value": {0: "batch_size"}, - }, - external_data=False, # Don't use external data format for simplicity; model should be small enough to fit in a single file - ) - print(f"AlphaZero model exported to ONNX at {path}") + # Subsequent saves: update only the weight tensors in the cached proto. + state_dict = self.evaluator.network.state_dict() + updated_count = 0 + for name, tensor in state_dict.items(): + if name in self._onnx_init_name_to_idx: + idx = self._onnx_init_name_to_idx[name] + self._onnx_proto.graph.initializer[idx].CopyFrom( + onnx.numpy_helper.from_array(tensor.cpu().numpy(), name=name) + ) + updated_count += 1 + expected_initializers = len(self._onnx_init_name_to_idx) + if updated_count != expected_initializers: + cached_names = set(self._onnx_init_name_to_idx.keys()) + missing_in_state_dict = sorted(cached_names - set(state_dict.keys())) + print( + "Warning: ONNX cached initializers not fully updated from state_dict " + f"({updated_count}/{expected_initializers} updated). " + f"Missing in state_dict: {missing_in_state_dict}" + ) + onnx.save(self._onnx_proto, str(path)) + print(f"AlphaZero model exported to ONNX at {path} (cached proto, weights updated)") def save_model_with_suffix(self, suffix: str) -> Path: path = resolve_path(self.params.model_dir, self.resolve_filename(suffix)) @@ -509,9 +532,7 @@ def end_game_batch(self, env=None): while replay_buffer and replay_buffer[-1]["value"] is None: position = replay_buffer.pop() agent = env.player_to_agent[position["player"]] - factor = 1 + self.params.game_length_bonus_factor * ( - 1 - env.game.completed_steps / env.max_steps - ) + factor = 1 + self.params.game_length_bonus_factor * (1 - env.game.completed_steps / env.max_steps) position["value"] = factor * env.rewards[agent] episode_positions.append(position) @@ -520,9 +541,7 @@ def end_game_batch(self, env=None): self.episode_count += 1 - def end_game_batch_and_save_replay_buffers( - self, temp_dir: Path, ready_dir: Path, model_version: int - ): + def end_game_batch_and_save_replay_buffers(self, temp_dir: Path, ready_dir: Path, model_version: int): if not self.params.training_mode: return @@ -536,9 +555,7 @@ def end_game_batch_and_save_replay_buffers( while replay_buffer and replay_buffer[-1]["value"] is None: position = replay_buffer.pop() agent = env.player_to_agent[position["player"]] - factor = 1 + self.params.game_length_bonus_factor * ( - 1 - env.game.completed_steps / env.max_steps - ) + factor = 1 + self.params.game_length_bonus_factor * (1 - env.game.completed_steps / env.max_steps) position["value"] = factor * env.rewards[agent] processed_replay_buffer.append(position) @@ -559,19 +576,11 @@ def end_game_batch_and_save_replay_buffers( npz_filename = temp_dir / f"{filename}.npz" np.savez_compressed( npz_filename, - input_arrays=np.stack( - [p["input_array"] for p in processed_replay_buffer] - ), + input_arrays=np.stack([p["input_array"] for p in processed_replay_buffer]), policies=np.stack([p["mcts_policy"] for p in processed_replay_buffer]), - action_masks=np.stack( - [p["action_mask"] for p in processed_replay_buffer] - ), - values=np.array( - [p["value"] for p in processed_replay_buffer], dtype=np.float32 - ), - players=np.array( - [int(p["player"]) for p in processed_replay_buffer], dtype=np.int32 - ), + action_masks=np.stack([p["action_mask"] for p in processed_replay_buffer]), + values=np.array([p["value"] for p in processed_replay_buffer], dtype=np.float32), + players=np.array([int(p["player"]) for p in processed_replay_buffer], dtype=np.int32), ) npz_filename.rename(ready_dir / npz_filename.name) @@ -584,10 +593,7 @@ def end_game(self, env): self.end_game_batch(env) # This is just used when we use the old train.py script - if ( - self.params.train_every is not None - and self.episode_count % self.params.train_every == 0 - ): + if self.params.train_every is not None and self.episode_count % self.params.train_every == 0: self.train_iteration() def compute_loss_and_reward(self, length: int) -> Tuple[float, float]: @@ -599,9 +605,7 @@ def compute_loss_and_reward(self, length: int) -> Tuple[float, float]: return float(avg_loss), 0.0 - def train_iteration( - self, is_replay_buffer_bootstrap=False, epoch=None, episode=None - ) -> bool: + def train_iteration(self, is_replay_buffer_bootstrap=False, epoch=None, episode=None) -> bool: """ Train the neural network on collected data. @@ -619,9 +623,7 @@ def log_loss_entry(entry): entry["loss_policy_val"], entry["loss_value_val"], ) - print( - f"{i:>7} {t:>7.3f} {vt:>7.3f} {p:>7.3f} {vp:>7.3f} {v:>7.3f} {vv:>7.3f}" - ) + print(f"{i:>7} {t:>7.3f} {vt:>7.3f} {p:>7.3f} {vp:>7.3f} {v:>7.3f} {vv:>7.3f}") else: print(f"{i:>7} {t:>7.3f} {p:>7.3f} {v:>7.3f}") @@ -633,10 +635,7 @@ def log_loss_entry(entry): del entry["completion"] self.wandb_run.log(entry) - if ( - len(self.replay_buffer) * (1.0 - self.params.validation_ratio) - < self.params.batch_size - ): + if len(self.replay_buffer) * (1.0 - self.params.validation_ratio) < self.params.batch_size: return False # Save replay buffer if requested @@ -650,15 +649,11 @@ def log_loss_entry(entry): self.first_replay_buffer_saved = True Timer.start("training") - print( - f"Training the network (buffer size: {len(self.replay_buffer)}, batch size: {self.params.batch_size})..." - ) + print(f"Training the network (buffer size: {len(self.replay_buffer)}, batch size: {self.params.batch_size})...") if self.params.validation_ratio > 0.0: print("==== Training & Validation Loss ====") - print( - " Epoch Total Val Total Policy Val Policy Value Val Value" - ) + print(" Epoch Total Val Total Policy Val Policy Value Val Value") else: print("==== Training Loss ====") print(" Epoch Total Policy Value") @@ -678,9 +673,7 @@ def _replay_buffer_filename(self, episode_number: int): params = { "ep": episode_number, "i": get_initial_random_seed(), - "t": int(self.params.temperature * 100) - if self.params.temperature - else None, + "t": int(self.params.temperature * 100) if self.params.temperature else None, "dt": self.params.drop_t_on_step, "rbs": self.params.replay_buffer_size, "n": self.params.mcts_n, @@ -712,9 +705,7 @@ def save_replay_buffer_to_file(self, episode_number: int): """Save replay buffer contents to a file.""" replay_buffer_dir = "replay_buffers" os.makedirs(replay_buffer_dir, exist_ok=True) - filepath = os.path.join( - replay_buffer_dir, self._replay_buffer_filename(episode_number) - ) + filepath = os.path.join(replay_buffer_dir, self._replay_buffer_filename(episode_number)) with open(filepath, "wb") as f: pickle.dump(list(self.replay_buffer), f) print(f"Saved replay buffer to {filepath}") @@ -724,12 +715,8 @@ def load_replay_buffer_from_file(self) -> bool: if self.params.train_every is None or not self.is_training(): return False - assert isinstance( - self.params.train_every, int - ), "train_every must be set to load replay buffer from file" - filepath = os.path.join( - "replay_buffers", self._replay_buffer_filename(self.params.train_every) - ) + assert isinstance(self.params.train_every, int), "train_every must be set to load replay buffer from file" + filepath = os.path.join("replay_buffers", self._replay_buffer_filename(self.params.train_every)) if not os.path.exists(filepath): return False @@ -757,9 +744,7 @@ def _log_action( self.action_log.action_score_ranking(scores) self.action_log.action_text(root_action, f"{root_value:0.2f}") - def get_action_batch( - self, observations_with_ids: list[tuple[int, dict]] - ) -> list[tuple[int, int]]: + def get_action_batch(self, observations_with_ids: list[tuple[int, dict]]) -> list[tuple[int, int]]: """ Get actions for multiple observations. Each entry in observation_with_ids needs a distinct id that can be an arbitrary number. The same id will be returned with the action matching the observation. @@ -801,17 +786,12 @@ def get_action_batch( ) temperature = self.initial_temperature - if ( - self.params.drop_t_on_step is not None - and game.completed_steps >= self.params.drop_t_on_step - ): + if self.params.drop_t_on_step is not None and game.completed_steps >= self.params.drop_t_on_step: temperature = 0 if temperature == 0.0: max_value = np.max(visit_probs) - visit_probs = np.array( - [1.0 if v == max_value else 0.0 for v in visit_probs] - ) + visit_probs = np.array([1.0 if v == max_value else 0.0 for v in visit_probs]) visit_probs /= np.sum(visit_probs) else: visit_probs = visit_probs ** (1.0 / temperature) @@ -830,13 +810,9 @@ def get_action_batch( # Store training data if in training mode if self.params.training_mode: # Convert visit counts to policy target (normalized) - policy_target = np.zeros( - self.action_encoder.num_actions, dtype=np.float32 - ) + policy_target = np.zeros(self.action_encoder.num_actions, dtype=np.float32) for child in root_children: - action_index = self.action_encoder.action_to_index( - child.action_taken - ) + action_index = self.action_encoder.action_to_index(child.action_taken) policy_target[action_index] = child.visit_count / visit_counts_sum self.store_training_data(game, policy_target, player, game_idx)