trustformers_optim/
onnx_export.rs

1//! ONNX Optimizer Export
2//!
3//! This module provides functionality to export optimizer configurations and states
4//! to ONNX format, enabling deployment and optimization in ONNX Runtime environments.
5
6use anyhow::{anyhow, Result};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10
11/// ONNX Graph Node representation for optimizer operations
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct ONNXNode {
14    pub name: String,
15    pub op_type: String,
16    pub inputs: Vec<String>,
17    pub outputs: Vec<String>,
18    pub attributes: HashMap<String, Value>,
19}
20
21/// ONNX Graph representation for optimizer
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct ONNXGraph {
24    pub name: String,
25    pub nodes: Vec<ONNXNode>,
26    pub inputs: Vec<String>,
27    pub outputs: Vec<String>,
28    pub initializers: HashMap<String, Vec<f32>>,
29}
30
31/// ONNX Model representation
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct ONNXModel {
34    pub ir_version: i64,
35    pub producer_name: String,
36    pub producer_version: String,
37    pub domain: String,
38    pub model_version: i64,
39    pub graph: ONNXGraph,
40}
41
42/// Optimizer configuration for ONNX export
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct OptimizerConfig {
45    pub optimizer_type: String,
46    pub learning_rate: f32,
47    pub parameters: HashMap<String, Value>,
48}
49
50/// ONNX Export configuration
51#[derive(Debug, Clone, Serialize, Deserialize)]
52pub struct ONNXExportConfig {
53    pub model_name: String,
54    pub opset_version: i64,
55    pub export_params: bool,
56    pub export_raw_ir: bool,
57    pub keep_initializers_as_inputs: bool,
58    pub custom_opsets: HashMap<String, i64>,
59    pub verbose: bool,
60}
61
62impl Default for ONNXExportConfig {
63    fn default() -> Self {
64        Self {
65            model_name: "TrustformeRS_Optimizer".to_string(),
66            opset_version: 17,
67            export_params: true,
68            export_raw_ir: false,
69            keep_initializers_as_inputs: false,
70            custom_opsets: HashMap::new(),
71            verbose: false,
72        }
73    }
74}
75
76/// ONNX Optimizer metadata for export
77#[derive(Debug, Clone, Serialize, Deserialize)]
78pub struct ONNXOptimizerMetadata {
79    pub optimizer_type: String,
80    pub version: String,
81    pub hyperparameters: HashMap<String, Value>,
82    pub state_variables: Vec<String>,
83    pub export_timestamp: String,
84    pub framework_version: String,
85}
86
87impl Default for ONNXOptimizerMetadata {
88    fn default() -> Self {
89        Self {
90            optimizer_type: "Adam".to_string(),
91            version: "1.0".to_string(),
92            hyperparameters: HashMap::new(),
93            state_variables: Vec::new(),
94            export_timestamp: "2025-07-22T00:00:00Z".to_string(),
95            framework_version: "0.1.0".to_string(),
96        }
97    }
98}
99
100/// ONNX Optimizer Exporter
101pub struct ONNXOptimizerExporter {
102    producer_name: String,
103    producer_version: String,
104}
105
106impl ONNXOptimizerExporter {
107    /// Create a new ONNX optimizer exporter
108    pub fn new() -> Self {
109        Self {
110            producer_name: "TrustformeRS".to_string(),
111            producer_version: "1.0.0".to_string(),
112        }
113    }
114
115    /// Export Adam optimizer to ONNX format
116    pub fn export_adam(
117        &self,
118        learning_rate: f32,
119        beta1: f32,
120        beta2: f32,
121        epsilon: f32,
122        weight_decay: f32,
123    ) -> Result<ONNXModel> {
124        let mut nodes = Vec::new();
125        let mut initializers = HashMap::new();
126
127        // Add learning rate as initializer
128        initializers.insert("learning_rate".to_string(), vec![learning_rate]);
129        initializers.insert("beta1".to_string(), vec![beta1]);
130        initializers.insert("beta2".to_string(), vec![beta2]);
131        initializers.insert("epsilon".to_string(), vec![epsilon]);
132        initializers.insert("weight_decay".to_string(), vec![weight_decay]);
133
134        // Create Adam optimizer node
135        let mut adam_attrs = HashMap::new();
136        adam_attrs.insert("alpha".to_string(), Value::from(learning_rate as f64));
137        adam_attrs.insert("beta".to_string(), Value::from(beta1 as f64));
138        adam_attrs.insert("beta2".to_string(), Value::from(beta2 as f64));
139        adam_attrs.insert("epsilon".to_string(), Value::from(epsilon as f64));
140        adam_attrs.insert("weight_decay".to_string(), Value::from(weight_decay as f64));
141
142        let adam_node = ONNXNode {
143            name: "adam_optimizer".to_string(),
144            op_type: "Adam".to_string(),
145            inputs: vec![
146                "gradients".to_string(),
147                "learning_rate".to_string(),
148                "beta1".to_string(),
149                "beta2".to_string(),
150                "epsilon".to_string(),
151                "weight_decay".to_string(),
152            ],
153            outputs: vec!["updated_parameters".to_string()],
154            attributes: adam_attrs,
155        };
156
157        nodes.push(adam_node);
158
159        let graph = ONNXGraph {
160            name: "adam_optimizer_graph".to_string(),
161            nodes,
162            inputs: vec!["gradients".to_string()],
163            outputs: vec!["updated_parameters".to_string()],
164            initializers,
165        };
166
167        Ok(ONNXModel {
168            ir_version: 7,
169            producer_name: self.producer_name.clone(),
170            producer_version: self.producer_version.clone(),
171            domain: "ai.onnx".to_string(),
172            model_version: 1,
173            graph,
174        })
175    }
176
177    /// Export SGD optimizer to ONNX format
178    pub fn export_sgd(
179        &self,
180        learning_rate: f32,
181        momentum: f32,
182        weight_decay: f32,
183        nesterov: bool,
184    ) -> Result<ONNXModel> {
185        let mut nodes = Vec::new();
186        let mut initializers = HashMap::new();
187
188        // Add hyperparameters as initializers
189        initializers.insert("learning_rate".to_string(), vec![learning_rate]);
190        initializers.insert("momentum".to_string(), vec![momentum]);
191        initializers.insert("weight_decay".to_string(), vec![weight_decay]);
192
193        // Create SGD optimizer node
194        let mut sgd_attrs = HashMap::new();
195        sgd_attrs.insert(
196            "learning_rate".to_string(),
197            Value::from(learning_rate as f64),
198        );
199        sgd_attrs.insert("momentum".to_string(), Value::from(momentum as f64));
200        sgd_attrs.insert("weight_decay".to_string(), Value::from(weight_decay as f64));
201        sgd_attrs.insert("nesterov".to_string(), Value::Bool(nesterov));
202
203        let sgd_node = ONNXNode {
204            name: "sgd_optimizer".to_string(),
205            op_type: "SGD".to_string(),
206            inputs: vec![
207                "gradients".to_string(),
208                "learning_rate".to_string(),
209                "momentum".to_string(),
210                "weight_decay".to_string(),
211            ],
212            outputs: vec!["updated_parameters".to_string()],
213            attributes: sgd_attrs,
214        };
215
216        nodes.push(sgd_node);
217
218        let graph = ONNXGraph {
219            name: "sgd_optimizer_graph".to_string(),
220            nodes,
221            inputs: vec!["gradients".to_string()],
222            outputs: vec!["updated_parameters".to_string()],
223            initializers,
224        };
225
226        Ok(ONNXModel {
227            ir_version: 7,
228            producer_name: self.producer_name.clone(),
229            producer_version: self.producer_version.clone(),
230            domain: "ai.onnx".to_string(),
231            model_version: 1,
232            graph,
233        })
234    }
235
236    /// Export AdamW optimizer to ONNX format
237    pub fn export_adamw(
238        &self,
239        learning_rate: f32,
240        beta1: f32,
241        beta2: f32,
242        epsilon: f32,
243        weight_decay: f32,
244    ) -> Result<ONNXModel> {
245        let mut nodes = Vec::new();
246        let mut initializers = HashMap::new();
247
248        // Add hyperparameters as initializers
249        initializers.insert("learning_rate".to_string(), vec![learning_rate]);
250        initializers.insert("beta1".to_string(), vec![beta1]);
251        initializers.insert("beta2".to_string(), vec![beta2]);
252        initializers.insert("epsilon".to_string(), vec![epsilon]);
253        initializers.insert("weight_decay".to_string(), vec![weight_decay]);
254
255        // Create AdamW optimizer node
256        let mut adamw_attrs = HashMap::new();
257        adamw_attrs.insert("alpha".to_string(), Value::from(learning_rate as f64));
258        adamw_attrs.insert("beta".to_string(), Value::from(beta1 as f64));
259        adamw_attrs.insert("beta2".to_string(), Value::from(beta2 as f64));
260        adamw_attrs.insert("epsilon".to_string(), Value::from(epsilon as f64));
261        adamw_attrs.insert("weight_decay".to_string(), Value::from(weight_decay as f64));
262
263        let adamw_node = ONNXNode {
264            name: "adamw_optimizer".to_string(),
265            op_type: "AdamW".to_string(),
266            inputs: vec![
267                "gradients".to_string(),
268                "learning_rate".to_string(),
269                "beta1".to_string(),
270                "beta2".to_string(),
271                "epsilon".to_string(),
272                "weight_decay".to_string(),
273            ],
274            outputs: vec!["updated_parameters".to_string()],
275            attributes: adamw_attrs,
276        };
277
278        nodes.push(adamw_node);
279
280        let graph = ONNXGraph {
281            name: "adamw_optimizer_graph".to_string(),
282            nodes,
283            inputs: vec!["gradients".to_string()],
284            outputs: vec!["updated_parameters".to_string()],
285            initializers,
286        };
287
288        Ok(ONNXModel {
289            ir_version: 7,
290            producer_name: self.producer_name.clone(),
291            producer_version: self.producer_version.clone(),
292            domain: "ai.onnx".to_string(),
293            model_version: 1,
294            graph,
295        })
296    }
297
298    /// Export optimizer configuration to JSON format for ONNX metadata
299    pub fn export_config(&self, config: &OptimizerConfig) -> Result<String> {
300        serde_json::to_string_pretty(config)
301            .map_err(|e| anyhow!("Failed to serialize optimizer config: {}", e))
302    }
303
304    /// Save ONNX model to file
305    pub fn save_model(&self, model: &ONNXModel, path: &str) -> Result<()> {
306        let json = serde_json::to_string_pretty(model)
307            .map_err(|e| anyhow!("Failed to serialize ONNX model: {}", e))?;
308
309        std::fs::write(path, json)
310            .map_err(|e| anyhow!("Failed to write ONNX model to file: {}", e))?;
311
312        Ok(())
313    }
314
315    /// Create optimizer config from common optimizers
316    pub fn create_adam_config(
317        &self,
318        learning_rate: f32,
319        beta1: f32,
320        beta2: f32,
321        epsilon: f32,
322        weight_decay: f32,
323    ) -> OptimizerConfig {
324        let mut parameters = HashMap::new();
325        parameters.insert("beta1".to_string(), Value::from(beta1 as f64));
326        parameters.insert("beta2".to_string(), Value::from(beta2 as f64));
327        parameters.insert("epsilon".to_string(), Value::from(epsilon as f64));
328        parameters.insert("weight_decay".to_string(), Value::from(weight_decay as f64));
329
330        OptimizerConfig {
331            optimizer_type: "Adam".to_string(),
332            learning_rate,
333            parameters,
334        }
335    }
336
337    pub fn create_sgd_config(
338        &self,
339        learning_rate: f32,
340        momentum: f32,
341        weight_decay: f32,
342        nesterov: bool,
343    ) -> OptimizerConfig {
344        let mut parameters = HashMap::new();
345        parameters.insert("momentum".to_string(), Value::from(momentum as f64));
346        parameters.insert("weight_decay".to_string(), Value::from(weight_decay as f64));
347        parameters.insert("nesterov".to_string(), Value::Bool(nesterov));
348
349        OptimizerConfig {
350            optimizer_type: "SGD".to_string(),
351            learning_rate,
352            parameters,
353        }
354    }
355
356    pub fn create_adamw_config(
357        &self,
358        learning_rate: f32,
359        beta1: f32,
360        beta2: f32,
361        epsilon: f32,
362        weight_decay: f32,
363    ) -> OptimizerConfig {
364        let mut parameters = HashMap::new();
365        parameters.insert("beta1".to_string(), Value::from(beta1 as f64));
366        parameters.insert("beta2".to_string(), Value::from(beta2 as f64));
367        parameters.insert("epsilon".to_string(), Value::from(epsilon as f64));
368        parameters.insert("weight_decay".to_string(), Value::from(weight_decay as f64));
369
370        OptimizerConfig {
371            optimizer_type: "AdamW".to_string(),
372            learning_rate,
373            parameters,
374        }
375    }
376}
377
378impl Default for ONNXOptimizerExporter {
379    fn default() -> Self {
380        Self::new()
381    }
382}
383
384/// Utility functions for ONNX export
385pub mod utils {
386    use super::*;
387
388    /// Validate ONNX model structure
389    pub fn validate_model(model: &ONNXModel) -> Result<()> {
390        if model.graph.nodes.is_empty() {
391            return Err(anyhow!("ONNX model must have at least one node"));
392        }
393
394        if model.graph.inputs.is_empty() {
395            return Err(anyhow!("ONNX model must have at least one input"));
396        }
397
398        if model.graph.outputs.is_empty() {
399            return Err(anyhow!("ONNX model must have at least one output"));
400        }
401
402        // Validate node connections
403        for node in &model.graph.nodes {
404            for input in &node.inputs {
405                if !model.graph.inputs.contains(input)
406                    && !model.graph.initializers.contains_key(input)
407                {
408                    // Check if input is output of another node
409                    let is_node_output =
410                        model.graph.nodes.iter().any(|n| n.outputs.contains(input));
411
412                    if !is_node_output {
413                        return Err(anyhow!("Node input '{}' is not connected", input));
414                    }
415                }
416            }
417        }
418
419        Ok(())
420    }
421
422    /// Create ONNX model with learning rate scheduler
423    pub fn create_with_scheduler(
424        optimizer_model: ONNXModel,
425        schedule_type: &str,
426        schedule_params: HashMap<String, f32>,
427    ) -> Result<ONNXModel> {
428        let mut model = optimizer_model;
429
430        // Add scheduler node
431        let mut scheduler_attrs = HashMap::new();
432        for (key, value) in schedule_params {
433            scheduler_attrs.insert(key, Value::from(value as f64));
434        }
435
436        let scheduler_node = ONNXNode {
437            name: "lr_scheduler".to_string(),
438            op_type: schedule_type.to_string(),
439            inputs: vec!["step".to_string()],
440            outputs: vec!["scheduled_learning_rate".to_string()],
441            attributes: scheduler_attrs,
442        };
443
444        model.graph.nodes.insert(0, scheduler_node);
445        model.graph.inputs.push("step".to_string());
446
447        // Update optimizer node to use scheduled learning rate
448        if let Some(optimizer_node) = model
449            .graph
450            .nodes
451            .iter_mut()
452            .find(|n| n.op_type == "Adam" || n.op_type == "SGD" || n.op_type == "AdamW")
453        {
454            if let Some(lr_input_idx) =
455                optimizer_node.inputs.iter().position(|i| i == "learning_rate")
456            {
457                optimizer_node.inputs[lr_input_idx] = "scheduled_learning_rate".to_string();
458            }
459        }
460
461        Ok(model)
462    }
463}
464
465#[cfg(test)]
466mod tests {
467    use super::*;
468
469    #[test]
470    fn test_onnx_adam_export() {
471        let exporter = ONNXOptimizerExporter::new();
472        let model = exporter
473            .export_adam(0.001, 0.9, 0.999, 1e-8, 0.01)
474            .expect("Operation failed in test");
475
476        assert_eq!(model.graph.name, "adam_optimizer_graph");
477        assert_eq!(model.graph.nodes.len(), 1);
478        assert_eq!(model.graph.nodes[0].op_type, "Adam");
479
480        utils::validate_model(&model).expect("Operation failed in test");
481    }
482
483    #[test]
484    fn test_onnx_sgd_export() {
485        let exporter = ONNXOptimizerExporter::new();
486        let model = exporter.export_sgd(0.01, 0.9, 1e-4, true).expect("Operation failed in test");
487
488        assert_eq!(model.graph.name, "sgd_optimizer_graph");
489        assert_eq!(model.graph.nodes.len(), 1);
490        assert_eq!(model.graph.nodes[0].op_type, "SGD");
491
492        utils::validate_model(&model).expect("Operation failed in test");
493    }
494
495    #[test]
496    fn test_onnx_adamw_export() {
497        let exporter = ONNXOptimizerExporter::new();
498        let model = exporter
499            .export_adamw(0.001, 0.9, 0.999, 1e-8, 0.01)
500            .expect("Operation failed in test");
501
502        assert_eq!(model.graph.name, "adamw_optimizer_graph");
503        assert_eq!(model.graph.nodes.len(), 1);
504        assert_eq!(model.graph.nodes[0].op_type, "AdamW");
505
506        utils::validate_model(&model).expect("Operation failed in test");
507    }
508
509    #[test]
510    fn test_config_creation() {
511        let exporter = ONNXOptimizerExporter::new();
512
513        let adam_config = exporter.create_adam_config(0.001, 0.9, 0.999, 1e-8, 0.01);
514        assert_eq!(adam_config.optimizer_type, "Adam");
515        assert_eq!(adam_config.learning_rate, 0.001);
516
517        let sgd_config = exporter.create_sgd_config(0.01, 0.9, 1e-4, true);
518        assert_eq!(sgd_config.optimizer_type, "SGD");
519        assert_eq!(sgd_config.learning_rate, 0.01);
520    }
521
522    #[test]
523    fn test_config_serialization() {
524        let exporter = ONNXOptimizerExporter::new();
525        let config = exporter.create_adam_config(0.001, 0.9, 0.999, 1e-8, 0.01);
526
527        let json = exporter.export_config(&config).expect("Operation failed in test");
528        assert!(json.contains("Adam"));
529        assert!(json.contains("0.001"));
530    }
531
532    #[test]
533    fn test_model_validation() {
534        let exporter = ONNXOptimizerExporter::new();
535        let model = exporter
536            .export_adam(0.001, 0.9, 0.999, 1e-8, 0.01)
537            .expect("Operation failed in test");
538
539        // Should pass validation
540        utils::validate_model(&model).expect("Operation failed in test");
541
542        // Test invalid model
543        let mut invalid_model = model.clone();
544        invalid_model.graph.nodes.clear();
545        assert!(utils::validate_model(&invalid_model).is_err());
546    }
547
548    #[test]
549    fn test_scheduler_integration() {
550        let exporter = ONNXOptimizerExporter::new();
551        let base_model = exporter
552            .export_adam(0.001, 0.9, 0.999, 1e-8, 0.01)
553            .expect("Operation failed in test");
554
555        let mut schedule_params = HashMap::new();
556        schedule_params.insert("decay_rate".to_string(), 0.95);
557
558        let model_with_scheduler =
559            utils::create_with_scheduler(base_model, "ExponentialDecay", schedule_params)
560                .expect("Operation failed in test");
561
562        assert_eq!(model_with_scheduler.graph.nodes.len(), 2);
563        assert_eq!(
564            model_with_scheduler.graph.nodes[0].op_type,
565            "ExponentialDecay"
566        );
567
568        utils::validate_model(&model_with_scheduler).expect("Operation failed in test");
569    }
570}
trustformers_optim/onnx_export.rs

trustformers_optim/
onnx_export.rs
