|
| 1 | +"""Example demonstrating save_model and model sharding functionality. |
| 2 | +
|
| 3 | +This example shows how to: |
| 4 | +1. Save an ONNX model with safetensors weights using save_model |
| 5 | +2. Shard large models across multiple safetensors files |
| 6 | +3. Load and verify sharded models with ONNX Runtime |
| 7 | +""" |
| 8 | + |
| 9 | +import glob |
| 10 | +import json |
| 11 | +import os |
| 12 | + |
| 13 | +import numpy as np |
| 14 | +import onnx |
| 15 | +import onnx.helper |
| 16 | +import onnx.numpy_helper |
| 17 | +import onnxruntime as ort |
| 18 | + |
| 19 | +import onnx_safetensors |
| 20 | + |
| 21 | + |
| 22 | +def create_example_model(large: bool = False) -> onnx.ModelProto: |
| 23 | + """Create an example ONNX model for demonstration. |
| 24 | +
|
| 25 | + Args: |
| 26 | + large: If True, creates a larger model to demonstrate sharding. |
| 27 | +
|
| 28 | + Returns: |
| 29 | + An ONNX model. |
| 30 | + """ |
| 31 | + if large: |
| 32 | + # Create a larger model with multiple weight matrices to demonstrate sharding |
| 33 | + weights1 = np.random.randn(1000, 1000).astype(np.float32) # ~4MB |
| 34 | + weights2 = np.random.randn(1000, 2000).astype(np.float32) # ~8MB |
| 35 | + weights3 = np.random.randn(2000, 1000).astype(np.float32) # ~8MB |
| 36 | + |
| 37 | + graph = onnx.helper.make_graph( |
| 38 | + [ |
| 39 | + onnx.helper.make_node("MatMul", ["input", "weights1"], ["temp1"]), |
| 40 | + onnx.helper.make_node("MatMul", ["temp1", "weights2"], ["temp2"]), |
| 41 | + onnx.helper.make_node("MatMul", ["temp2", "weights3"], ["output"]), |
| 42 | + ], |
| 43 | + "large_model", |
| 44 | + inputs=[ |
| 45 | + onnx.helper.make_tensor_value_info( |
| 46 | + "input", onnx.TensorProto.FLOAT, [1, 1000] |
| 47 | + ), |
| 48 | + ], |
| 49 | + outputs=[ |
| 50 | + onnx.helper.make_tensor_value_info( |
| 51 | + "output", onnx.TensorProto.FLOAT, [1, 1000] |
| 52 | + ), |
| 53 | + ], |
| 54 | + initializer=[ |
| 55 | + onnx.numpy_helper.from_array(weights1, name="weights1"), |
| 56 | + onnx.numpy_helper.from_array(weights2, name="weights2"), |
| 57 | + onnx.numpy_helper.from_array(weights3, name="weights3"), |
| 58 | + ], |
| 59 | + ) |
| 60 | + else: |
| 61 | + # Create a simple model |
| 62 | + weights = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32) |
| 63 | + |
| 64 | + graph = onnx.helper.make_graph( |
| 65 | + [ |
| 66 | + onnx.helper.make_node("Add", ["input", "weights"], ["output"]), |
| 67 | + ], |
| 68 | + "simple_model", |
| 69 | + inputs=[ |
| 70 | + onnx.helper.make_tensor_value_info( |
| 71 | + "input", onnx.TensorProto.FLOAT, [2, 3] |
| 72 | + ), |
| 73 | + ], |
| 74 | + outputs=[ |
| 75 | + onnx.helper.make_tensor_value_info( |
| 76 | + "output", onnx.TensorProto.FLOAT, [2, 3] |
| 77 | + ), |
| 78 | + ], |
| 79 | + initializer=[onnx.numpy_helper.from_array(weights, name="weights")], |
| 80 | + ) |
| 81 | + |
| 82 | + model = onnx.helper.make_model( |
| 83 | + graph, opset_imports=[onnx.helper.make_opsetid("", 14)], ir_version=10 |
| 84 | + ) |
| 85 | + return model |
| 86 | + |
| 87 | + |
| 88 | +def example_basic_save_model(): |
| 89 | + """Example 1: Basic usage of save_model.""" |
| 90 | + print("Example 1: Basic save_model usage") |
| 91 | + print("=" * 50) |
| 92 | + |
| 93 | + # Create a simple model |
| 94 | + model = create_example_model(large=False) |
| 95 | + |
| 96 | + # Save model and weights |
| 97 | + # This creates: |
| 98 | + # - simple_model.onnx (ONNX model file) |
| 99 | + # - simple_model.safetensors (weights file) |
| 100 | + onnx_safetensors.save_model(model, "simple_model.onnx") |
| 101 | + |
| 102 | + print("✓ Saved simple_model.onnx and simple_model.safetensors") |
| 103 | + |
| 104 | + # Load and verify the model with ONNX Runtime |
| 105 | + sess = ort.InferenceSession("simple_model.onnx", providers=["CPUExecutionProvider"]) |
| 106 | + input_data = np.ones((2, 3), dtype=np.float32) |
| 107 | + outputs = sess.run(None, {"input": input_data}) |
| 108 | + |
| 109 | + print("✓ Model runs successfully with ONNX Runtime") |
| 110 | + print(f" Output shape: {outputs[0].shape}") |
| 111 | + print() |
| 112 | + |
| 113 | + |
| 114 | +def example_custom_weights_file(): |
| 115 | + """Example 2: Specify a custom name for the weights file.""" |
| 116 | + print("Example 2: Custom weights file name") |
| 117 | + print("=" * 50) |
| 118 | + |
| 119 | + model = create_example_model(large=False) |
| 120 | + |
| 121 | + # Save with custom weights file name |
| 122 | + # This creates: |
| 123 | + # - my_model.onnx |
| 124 | + # - custom_weights.safetensors |
| 125 | + onnx_safetensors.save_model( |
| 126 | + model, "my_model.onnx", external_data="custom_weights.safetensors" |
| 127 | + ) |
| 128 | + |
| 129 | + print("✓ Saved my_model.onnx with custom_weights.safetensors") |
| 130 | + print() |
| 131 | + |
| 132 | + |
| 133 | +def example_model_sharding(): |
| 134 | + """Example 3: Shard a large model across multiple files.""" |
| 135 | + print("Example 3: Model sharding") |
| 136 | + print("=" * 50) |
| 137 | + |
| 138 | + # Create a larger model |
| 139 | + model = create_example_model(large=True) |
| 140 | + |
| 141 | + # Shard the model with 5MB per shard |
| 142 | + # This creates: |
| 143 | + # - large_model.onnx |
| 144 | + # - large_model-00001-of-00004.safetensors |
| 145 | + # - large_model-00002-of-00004.safetensors |
| 146 | + # - large_model-00003-of-00004.safetensors |
| 147 | + # - large_model-00004-of-00004.safetensors |
| 148 | + # - large_model.safetensors.index.json (index file) |
| 149 | + onnx_safetensors.save_model(model, "large_model.onnx", max_shard_size="5MB") |
| 150 | + |
| 151 | + print("✓ Saved large_model.onnx with sharded weights") |
| 152 | + print(" Files created:") |
| 153 | + |
| 154 | + # List the created shard files |
| 155 | + shard_files = sorted(glob.glob("large_model-*.safetensors")) |
| 156 | + for shard_file in shard_files: |
| 157 | + size_mb = os.path.getsize(shard_file) / (1024 * 1024) |
| 158 | + print(f" - {shard_file} ({size_mb:.2f} MB)") |
| 159 | + |
| 160 | + # Check for index file |
| 161 | + if os.path.exists("large_model.safetensors.index.json"): |
| 162 | + with open("large_model.safetensors.index.json") as f: |
| 163 | + index = json.load(f) |
| 164 | + print(f" ✓ Index file created with {len(index['weight_map'])} tensors mapped") |
| 165 | + |
| 166 | + # Verify the sharded model works with ONNX Runtime |
| 167 | + sess = ort.InferenceSession("large_model.onnx", providers=["CPUExecutionProvider"]) |
| 168 | + input_data = np.random.randn(1, 1000).astype(np.float32) |
| 169 | + outputs = sess.run(None, {"input": input_data}) |
| 170 | + |
| 171 | + print("✓ Sharded model runs successfully with ONNX Runtime") |
| 172 | + print(f" Output shape: {outputs[0].shape}") |
| 173 | + print() |
| 174 | + |
| 175 | + |
| 176 | +def example_save_file_with_sharding(): |
| 177 | + """Example 4: Use save_file with sharding for more control.""" |
| 178 | + print("Example 4: save_file with sharding") |
| 179 | + print("=" * 50) |
| 180 | + |
| 181 | + model = create_example_model(large=True) |
| 182 | + |
| 183 | + # Save only the weights with sharding |
| 184 | + # Note: This doesn't save the ONNX model file itself |
| 185 | + onnx_safetensors.save_file( |
| 186 | + model, |
| 187 | + "weights_only.safetensors", |
| 188 | + base_dir=".", |
| 189 | + max_shard_size="5MB", |
| 190 | + replace_data=False, # Don't modify the model |
| 191 | + ) |
| 192 | + |
| 193 | + print("✓ Saved sharded weights without modifying the model") |
| 194 | + |
| 195 | + shard_files = sorted(glob.glob("weights_only-*.safetensors")) |
| 196 | + print(f" Created {len(shard_files)} shard files") |
| 197 | + print() |
| 198 | + |
| 199 | + |
| 200 | +if __name__ == "__main__": |
| 201 | + print("ONNX-Safetensors: save_model and Sharding Examples") |
| 202 | + print("=" * 50) |
| 203 | + print() |
| 204 | + |
| 205 | + # Run all examples |
| 206 | + example_basic_save_model() |
| 207 | + example_custom_weights_file() |
| 208 | + example_model_sharding() |
| 209 | + example_save_file_with_sharding() |
| 210 | + |
| 211 | + print("All examples completed successfully! ✓") |
| 212 | + print() |
| 213 | + print("Note: This example created several files for demonstration.") |
| 214 | + print("You can safely delete them after reviewing.") |
0 commit comments