Merge pull request #341 from adamantivm/jac/sample_onnx_resnet_model

Updates ONNX save code and experiment to support ResNet

Author: adamantivm

deep_quoridor/MODEL_SAVE_OPTIONS.md

@@ -2,7 +2,7 @@
 
 ## Summary
 
-Added configuration options to save AlphaZero models in PyTorch and/or ONNX formats during training.
+Added configuration options to save AlphaZero models in ONNX format during training. PyTorch format (.pt files) is always saved.
 
 ## Configuration Options
 
@@ -12,13 +12,12 @@ Add these parameters to the `training` section of your YAML configuration:
 training:
   # ... other training parameters ...
   model_save_timing: false  # Set to true to print timing information for model saving
-  save_pytorch: true        # Save models in PyTorch format (.pt files) - DEFAULT
   save_onnx: false          # Save models in ONNX format (.onnx files) - DEFAULT
 ```
 
 ## Default Behavior
 
-- **PyTorch format**: Enabled by default (`save_pytorch: true`)
+- **PyTorch format**: Always enabled (cannot be disabled)
 - **ONNX format**: Disabled by default (`save_onnx: false`)
 - **Timing output**: Disabled by default (`model_save_timing: false`)
 
@@ -29,31 +28,22 @@ This ensures backward compatibility with existing configurations.
 ### Example 1: Save only PyTorch format (default)
 ```yaml
 training:
-  save_pytorch: true
   save_onnx: false
 ```
 
 ### Example 2: Save both PyTorch and ONNX formats
 ```yaml
 training:
-  save_pytorch: true
   save_onnx: true
   model_save_timing: true  # See timing for both formats
 ```
 
-### Example 3: Save only ONNX format
-```yaml
-training:
-  save_pytorch: false
-  save_onnx: true
-```
-
 ## Test Configurations
 
 Two test configurations are provided:
 
 1. **`experiments/test_model_save_timing.yaml`** - Basic test with PyTorch only
-2. **`experiments/test_onnx_export.yaml`** - Test with both PyTorch and ONNX export enabled
+2. **`experiments/test_onnx_export.yaml`** - Test with ONNX export enabled
 
 ## ONNX Export Details
 
@@ -67,11 +57,11 @@ The ONNX export includes:
 ## Files Modified
 
 1. **`src/v2/config.py`**
-   - Added `save_pytorch`, `save_onnx`, and `model_save_timing` to `TrainingConfig`
+   - Added `save_onnx` and `model_save_timing` to `TrainingConfig`
 
 2. **`src/v2/trainer.py`**
-   - Updated initial model save (model_0) to support both formats
-   - Updated training loop model saves to support both formats
+   - Updated initial model save (model_0) to support ONNX format
+   - Updated training loop model saves to support ONNX format
    - Enhanced timing output to show which formats were saved
 
 3. **`src/agents/alphazero/alphazero.py`**

deep_quoridor/experiments/B5W3/test_onnx_export.yaml

@@ -0,0 +1,16 @@
+extend: base.yaml
+run_id: sample-onnx-export-$DATETIME
+
+alphazero:
+  mcts_n: 50
+wandb: null  # Disable wandb for testing
+self_play:
+  num_workers: 1
+  parallel_games: 2
+  alphazero:
+    mcts_noise_epsilon: 0.25
+training:
+  finish_after: 10 models
+  model_save_timing: true 
+  save_onnx: true
+benchmarks: []

deep_quoridor/experiments/test_model_save_timing.yaml

@@ -419,17 +419,24 @@ def save_model(self, path):
     def save_model_onnx(self, path):
         """Export the model to ONNX format."""
         import torch.onnx
-        
+
         # Create directory for saving models if it doesn't exist
         os.makedirs(Path(path).absolute().parents[0], exist_ok=True)
-        
+
         # Set the network to evaluation mode
         self.evaluator.network.eval()
-        
+
         # Create a dummy input tensor with the correct shape
-        # The input size is determined by the network's input_size attribute
-        dummy_input = torch.randn(1, self.evaluator.network.input_size, device=self.device)
-        
+        # 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)
+        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,
@@ -445,6 +452,7 @@ def save_model_onnx(self, path):
                 "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}")
 

deep_quoridor/experiments/test_onnx_export.yaml

@@ -76,7 +76,6 @@ class TrainingConfig(StrictBaseModel):
     weight_decay: float
     replay_buffer_size: int
     model_save_timing: bool = False
-    save_pytorch: bool = True
     save_onnx: bool = False
     finish_after: Optional[str] = None
 

deep_quoridor/src/agents/alphazero/alphazero.py

@@ -64,11 +64,10 @@ def train(config: Config):
         wandb_run = MockWandb()
 
     # Save initial model (model_0)
-    if config.training.save_pytorch:
-        filename = config.paths.checkpoints / "model_0.pt"
-        alphazero_agent.save_model(filename)
-        LatestModel.write(config, str(filename), 0)
-    
+    filename = config.paths.checkpoints / "model_0.pt"
+    alphazero_agent.save_model(filename)
+    LatestModel.write(config, str(filename), 0)
+
     if config.training.save_onnx:
         onnx_filename = config.paths.checkpoints / "model_0.onnx"
         alphazero_agent.save_model_onnx(onnx_filename)
@@ -166,29 +165,27 @@ def train(config: Config):
         )
 
         Timer.start("save-model")
-        
-        # Save in PyTorch format if enabled
-        if config.training.save_pytorch:
-            new_model_filename = config.paths.checkpoints / f"model_{model_version}.pt"
-            alphazero_agent.save_model(new_model_filename)
-            LatestModel.write(config, str(new_model_filename), model_version)
-        
+
+        # Save in PyTorch format
+        new_model_filename = config.paths.checkpoints / f"model_{model_version}.pt"
+        alphazero_agent.save_model(new_model_filename)
+        LatestModel.write(config, str(new_model_filename), model_version)
+
         # Save in ONNX format if enabled
         if config.training.save_onnx:
             onnx_model_filename = config.paths.checkpoints / f"model_{model_version}.onnx"
             alphazero_agent.save_model_onnx(onnx_model_filename)
-        
+
         time_save_model = Timer.finish("save-model")
-        
+
         if config.training.model_save_timing:
             formats = []
-            if config.training.save_pytorch:
-                formats.append("PyTorch")
+            formats.append("PyTorch")
             if config.training.save_onnx:
                 formats.append("ONNX")
             format_str = " and ".join(formats) if formats else "no format"
             print(f"Saving model ({format_str}) took {time_save_model:.4f}s")
-        
+
         model_version += 1
 
     ShutdownSignal.signal(config)

deep_quoridor/src/v2/config.py

@@ -51,11 +51,6 @@ def test_override_boolean_true(config_file):
     assert config.training.model_save_timing is True
 
 
-def test_override_boolean_false(config_file):
-    config = load_user_config(config_file, overrides=["training.save_pytorch=false"])
-    assert config.training.save_pytorch is False
-
-
 def test_override_int(config_file):
     config = load_user_config(config_file, overrides=["alphazero.mcts_n=500"])
     assert config.alphazero.mcts_n == 500