torsh_nn/
summary.rs

1//! Model summary utilities for analyzing neural network architectures
2//!
3//! This module provides tools for printing detailed summaries of neural network models,
4//! including layer information, parameter counts, and memory usage estimates.
5
6use crate::{Module, Parameter};
7use torsh_core::error::Result;
8use torsh_tensor::Tensor;
9
10// Conditional imports for std/no_std compatibility
11#[cfg(feature = "std")]
12use std::{
13    collections::HashMap,
14    fmt::{self, Display},
15    string::String,
16    time::{Duration, Instant},
17    vec::Vec,
18};
19
20#[cfg(not(feature = "std"))]
21use alloc::{
22    fmt::{self, Display},
23    string::String,
24    vec::Vec,
25};
26
27#[cfg(not(feature = "std"))]
28use hashbrown::HashMap;
29
30/// Information about a single layer in the model
31#[derive(Debug, Clone)]
32pub struct LayerInfo {
33    /// Layer name/identifier
34    pub name: String,
35    /// Layer type (e.g., "Linear", "Conv2d", "ReLU")
36    pub layer_type: String,
37    /// Input shape
38    pub input_shape: Vec<usize>,
39    /// Output shape
40    pub output_shape: Vec<usize>,
41    /// Number of parameters
42    pub param_count: usize,
43    /// Number of trainable parameters
44    pub trainable_params: usize,
45    /// Memory usage estimate in bytes
46    pub memory_bytes: usize,
47}
48
49impl LayerInfo {
50    /// Create a new LayerInfo
51    pub fn new(
52        name: String,
53        layer_type: String,
54        input_shape: Vec<usize>,
55        output_shape: Vec<usize>,
56        param_count: usize,
57        trainable_params: usize,
58    ) -> Self {
59        // Estimate memory usage (rough approximation)
60        let input_elements: usize = input_shape.iter().product();
61        let output_elements: usize = output_shape.iter().product();
62        let memory_bytes = (input_elements + output_elements + param_count) * 4; // Assuming f32
63
64        Self {
65            name,
66            layer_type,
67            input_shape,
68            output_shape,
69            param_count,
70            trainable_params,
71            memory_bytes,
72        }
73    }
74}
75
76impl Display for LayerInfo {
77    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
78        write!(
79            f,
80            "{:<20} {:<15} {:<20} {:<20} {:>10} {:>10}",
81            self.name,
82            self.layer_type,
83            format!("{:?}", self.input_shape),
84            format!("{:?}", self.output_shape),
85            format_number(self.param_count),
86            format_bytes(self.memory_bytes)
87        )
88    }
89}
90
91/// Complete model summary information
92#[derive(Debug, Clone)]
93pub struct ModelSummary {
94    /// Information for each layer
95    pub layers: Vec<LayerInfo>,
96    /// Total number of parameters
97    pub total_params: usize,
98    /// Total number of trainable parameters
99    pub trainable_params: usize,
100    /// Total memory usage estimate in bytes
101    pub total_memory_bytes: usize,
102    /// Model input shape
103    pub input_shape: Vec<usize>,
104    /// Model output shape
105    pub output_shape: Vec<usize>,
106}
107
108impl ModelSummary {
109    /// Create a new model summary
110    pub fn new(layers: Vec<LayerInfo>, input_shape: Vec<usize>, output_shape: Vec<usize>) -> Self {
111        let total_params = layers.iter().map(|l| l.param_count).sum();
112        let trainable_params = layers.iter().map(|l| l.trainable_params).sum();
113        let total_memory_bytes = layers.iter().map(|l| l.memory_bytes).sum();
114
115        Self {
116            layers,
117            total_params,
118            trainable_params,
119            total_memory_bytes,
120            input_shape,
121            output_shape,
122        }
123    }
124
125    /// Print a formatted summary to stdout
126    pub fn print(&self) {
127        println!("{}", self);
128    }
129}
130
131impl Display for ModelSummary {
132    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
133        writeln!(f, "========================================================================================")?;
134        writeln!(f, "Model Summary")?;
135        writeln!(f, "========================================================================================")?;
136        writeln!(f, "Input Shape: {:?}", self.input_shape)?;
137        writeln!(f, "Output Shape: {:?}", self.output_shape)?;
138        writeln!(f, "========================================================================================")?;
139        writeln!(
140            f,
141            "{:<20} {:<15} {:<20} {:<20} {:>10} {:>10}",
142            "Layer (type)", "Type", "Input Shape", "Output Shape", "Param #", "Memory"
143        )?;
144        writeln!(f, "========================================================================================")?;
145
146        for layer in &self.layers {
147            writeln!(f, "{}", layer)?;
148        }
149
150        writeln!(f, "========================================================================================")?;
151        writeln!(f, "Total params: {}", format_number(self.total_params))?;
152        writeln!(
153            f,
154            "Trainable params: {}",
155            format_number(self.trainable_params)
156        )?;
157        writeln!(
158            f,
159            "Non-trainable params: {}",
160            format_number(self.total_params - self.trainable_params)
161        )?;
162        writeln!(
163            f,
164            "Total memory usage: {}",
165            format_bytes(self.total_memory_bytes)
166        )?;
167        writeln!(f, "========================================================================================")?;
168
169        Ok(())
170    }
171}
172
173/// Summary configuration options
174#[derive(Debug, Clone)]
175pub struct SummaryConfig {
176    /// Maximum depth to traverse in nested modules
177    pub max_depth: usize,
178    /// Whether to show only trainable parameters
179    pub trainable_only: bool,
180    /// Whether to include memory estimates
181    pub show_memory: bool,
182    /// Whether to use verbose output
183    pub verbose: bool,
184}
185
186impl Default for SummaryConfig {
187    fn default() -> Self {
188        Self {
189            max_depth: 10,
190            trainable_only: false,
191            show_memory: true,
192            verbose: false,
193        }
194    }
195}
196
197/// Create a summary of a model
198pub fn summarize<M: Module>(
199    model: &M,
200    input_shape: &[usize],
201    config: Option<SummaryConfig>,
202) -> Result<ModelSummary> {
203    let config = config.unwrap_or_default();
204
205    // Create a dummy input tensor to trace through the model
206    let dummy_input = torsh_tensor::creation::zeros(input_shape)?;
207
208    // Get model output to determine output shape
209    let output = model.forward(&dummy_input)?;
210    let output_shape = output.shape().dims().to_vec();
211
212    // Analyze the model structure
213    let layers = analyze_model_structure(model, input_shape, &config)?;
214
215    Ok(ModelSummary::new(
216        layers,
217        input_shape.to_vec(),
218        output_shape,
219    ))
220}
221
222/// Analyze the structure of a model and extract layer information
223fn analyze_model_structure<M: Module>(
224    model: &M,
225    input_shape: &[usize],
226    config: &SummaryConfig,
227) -> Result<Vec<LayerInfo>> {
228    let mut layers = Vec::new();
229
230    // Get all parameters
231    let parameters = model.parameters();
232    let _named_parameters = model.named_parameters();
233
234    // For simplicity, we'll create a single layer info for the entire model
235    // In a more sophisticated implementation, we would traverse the module tree
236    let total_params = count_parameters(&parameters);
237    let trainable_params = count_trainable_parameters(&parameters);
238
239    let layer_info = LayerInfo::new(
240        "Model".to_string(),
241        get_module_type_name(model),
242        input_shape.to_vec(),
243        input_shape.to_vec(), // Placeholder - would be computed from forward pass
244        total_params,
245        trainable_params,
246    );
247
248    layers.push(layer_info);
249
250    // If the model has children, analyze them recursively
251    if config.max_depth > 0 {
252        let children = model.children();
253        for (i, child) in children.iter().enumerate() {
254            let _child_config = SummaryConfig {
255                max_depth: config.max_depth - 1,
256                ..config.clone()
257            };
258
259            let child_name = format!("child_{}", i);
260            let child_params = child.parameters();
261            let child_total_params = count_parameters(&child_params);
262            let child_trainable_params = count_trainable_parameters(&child_params);
263
264            let child_info = LayerInfo::new(
265                child_name,
266                get_module_type_name(*child),
267                input_shape.to_vec(), // Simplified
268                input_shape.to_vec(), // Simplified
269                child_total_params,
270                child_trainable_params,
271            );
272
273            layers.push(child_info);
274        }
275    }
276
277    Ok(layers)
278}
279
280/// Count total parameters in a parameter map
281fn count_parameters(parameters: &HashMap<String, Parameter>) -> usize {
282    parameters
283        .values()
284        .map(|param| {
285            let tensor_guard = param.tensor();
286            let tensor = tensor_guard.read();
287            tensor.shape().dims().iter().product::<usize>()
288        })
289        .sum()
290}
291
292/// Count trainable parameters in a parameter map
293fn count_trainable_parameters(parameters: &HashMap<String, Parameter>) -> usize {
294    // For now, assume all parameters are trainable
295    // In a full implementation, this would check the requires_grad flag
296    count_parameters(parameters)
297}
298
299/// Get the type name of a module (simplified implementation)
300fn get_module_type_name<M: Module + ?Sized>(_module: &M) -> String {
301    // This is a simplified implementation
302    // In a full implementation, we would use type reflection or naming conventions
303    "Module".to_string()
304}
305
306/// Format a number with appropriate units (K, M, B)
307fn format_number(num: usize) -> String {
308    if num >= 1_000_000_000 {
309        format!("{:.1}B", num as f64 / 1_000_000_000.0)
310    } else if num >= 1_000_000 {
311        format!("{:.1}M", num as f64 / 1_000_000.0)
312    } else if num >= 1_000 {
313        format!("{:.1}K", num as f64 / 1_000.0)
314    } else {
315        num.to_string()
316    }
317}
318
319/// Format bytes with appropriate units (KB, MB, GB)
320fn format_bytes(bytes: usize) -> String {
321    if bytes >= 1_073_741_824 {
322        format!("{:.1} GB", bytes as f64 / 1_073_741_824.0)
323    } else if bytes >= 1_048_576 {
324        format!("{:.1} MB", bytes as f64 / 1_048_576.0)
325    } else if bytes >= 1_024 {
326        format!("{:.1} KB", bytes as f64 / 1_024.0)
327    } else {
328        format!("{} B", bytes)
329    }
330}
331
332/// Estimate the memory usage of a tensor shape
333pub fn estimate_tensor_memory(shape: &[usize], dtype_size: usize) -> usize {
334    shape.iter().product::<usize>() * dtype_size
335}
336
337/// Advanced model profiler that can track memory usage and compute statistics
338pub struct ModelProfiler {
339    /// Whether to track memory usage
340    pub track_memory: bool,
341    /// Whether to track computation time
342    pub track_time: bool,
343    /// Whether to track activations
344    pub track_activations: bool,
345}
346
347impl Default for ModelProfiler {
348    fn default() -> Self {
349        Self {
350            track_memory: true,
351            track_time: false,
352            track_activations: false,
353        }
354    }
355}
356
357impl ModelProfiler {
358    /// Create a new model profiler
359    pub fn new() -> Self {
360        Self::default()
361    }
362
363    /// Enable memory tracking
364    pub fn with_memory_tracking(mut self) -> Self {
365        self.track_memory = true;
366        self
367    }
368
369    /// Enable time tracking
370    pub fn with_time_tracking(mut self) -> Self {
371        self.track_time = true;
372        self
373    }
374
375    /// Enable activation tracking
376    pub fn with_activation_tracking(mut self) -> Self {
377        self.track_activations = true;
378        self
379    }
380
381    /// Profile a model with the given input
382    pub fn profile<M: Module>(&self, model: &M, input: &Tensor) -> Result<ProfileResult> {
383        let start_memory = if self.track_memory {
384            Some(get_memory_usage())
385        } else {
386            None
387        };
388
389        let start_time = if self.track_time {
390            Some(std::time::Instant::now())
391        } else {
392            None
393        };
394
395        // Run forward pass
396        let output = model.forward(input)?;
397
398        let end_time = start_time.map(|start| start.elapsed());
399        let memory_used = start_memory.map(|start| get_memory_usage() - start);
400
401        Ok(ProfileResult {
402            input_shape: input.shape().dims().to_vec(),
403            output_shape: output.shape().dims().to_vec(),
404            memory_used,
405            execution_time: end_time,
406            parameter_count: count_parameters(&model.parameters()),
407        })
408    }
409}
410
411/// Result of model profiling
412#[derive(Debug, Clone)]
413pub struct ProfileResult {
414    /// Input tensor shape
415    pub input_shape: Vec<usize>,
416    /// Output tensor shape
417    pub output_shape: Vec<usize>,
418    /// Memory used during forward pass (if tracked)
419    pub memory_used: Option<usize>,
420    /// Execution time (if tracked)
421    pub execution_time: Option<std::time::Duration>,
422    /// Total parameter count
423    pub parameter_count: usize,
424}
425
426impl Display for ProfileResult {
427    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
428        writeln!(f, "Profile Result:")?;
429        writeln!(f, "  Input Shape: {:?}", self.input_shape)?;
430        writeln!(f, "  Output Shape: {:?}", self.output_shape)?;
431        writeln!(f, "  Parameters: {}", format_number(self.parameter_count))?;
432
433        if let Some(memory) = self.memory_used {
434            writeln!(f, "  Memory Used: {}", format_bytes(memory))?;
435        }
436
437        if let Some(time) = self.execution_time {
438            writeln!(f, "  Execution Time: {:.3}ms", time.as_secs_f64() * 1000.0)?;
439        }
440
441        Ok(())
442    }
443}
444
445/// Get current memory usage (simplified implementation)
446fn get_memory_usage() -> usize {
447    // This is a placeholder implementation
448    // In practice, you would use system-specific APIs to get actual memory usage
449    0
450}
451
452/// Utility functions for quick model analysis
453pub mod utils {
454    use super::*;
455
456    /// Quick summary with default configuration
457    pub fn quick_summary<M: Module>(model: &M, input_shape: &[usize]) -> Result<()> {
458        let summary = summarize(model, input_shape, None)?;
459        summary.print();
460        Ok(())
461    }
462
463    /// Count total parameters in a model
464    pub fn count_model_parameters<M: Module>(model: &M) -> usize {
465        count_parameters(&model.parameters())
466    }
467
468    /// Get model size in MB (assuming f32 parameters)
469    pub fn get_model_size_mb<M: Module>(model: &M) -> f64 {
470        let param_count = count_model_parameters(model);
471        (param_count * 4) as f64 / 1_048_576.0 // 4 bytes per f32, convert to MB
472    }
473
474    /// Check if model fits in given memory budget
475    pub fn check_memory_budget<M: Module>(
476        model: &M,
477        input_shape: &[usize],
478        budget_mb: f64,
479    ) -> bool {
480        let model_size = get_model_size_mb(model);
481        let input_size = estimate_tensor_memory(input_shape, 4) as f64 / 1_048_576.0;
482        let estimated_total = model_size + input_size * 2.0; // Factor for intermediate activations
483
484        estimated_total <= budget_mb
485    }
486}
487
488/// Enhanced profiling tools for comprehensive model analysis
489pub mod profiling {
490    use super::*;
491
492    /// FLOPS counter for different layer types
493    #[derive(Debug, Clone)]
494    pub struct FLOPSCounter {
495        pub total_flops: u64,
496        pub layer_flops: HashMap<String, u64>,
497    }
498
499    impl FLOPSCounter {
500        pub fn new() -> Self {
501            Self {
502                total_flops: 0,
503                layer_flops: HashMap::new(),
504            }
505        }
506
507        /// Estimate FLOPS for a linear layer
508        pub fn count_linear_flops(
509            &mut self,
510            layer_name: String,
511            input_size: usize,
512            output_size: usize,
513            batch_size: usize,
514        ) {
515            let flops = (batch_size * input_size * output_size * 2) as u64; // 2 for multiply-add
516            self.layer_flops.insert(layer_name, flops);
517            self.total_flops += flops;
518        }
519
520        /// Estimate FLOPS for a convolution layer
521        pub fn count_conv_flops(
522            &mut self,
523            layer_name: String,
524            input_shape: &[usize],
525            kernel_size: &[usize],
526            output_channels: usize,
527        ) {
528            let batch_size = input_shape[0];
529            let input_channels = input_shape[1];
530            let output_height = input_shape[2]; // Simplified
531            let output_width = input_shape[3]; // Simplified
532
533            let kernel_flops = kernel_size.iter().product::<usize>() * input_channels;
534            let output_pixels = output_height * output_width * output_channels;
535            let flops = (batch_size * output_pixels * kernel_flops * 2) as u64; // 2 for multiply-add
536
537            self.layer_flops.insert(layer_name, flops);
538            self.total_flops += flops;
539        }
540
541        /// Get formatted FLOPS string
542        pub fn format_flops(flops: u64) -> String {
543            if flops >= 1_000_000_000_000 {
544                format!("{:.2} TFLOPS", flops as f64 / 1_000_000_000_000.0)
545            } else if flops >= 1_000_000_000 {
546                format!("{:.2} GFLOPS", flops as f64 / 1_000_000_000.0)
547            } else if flops >= 1_000_000 {
548                format!("{:.2} MFLOPS", flops as f64 / 1_000_000.0)
549            } else if flops >= 1_000 {
550                format!("{:.2} KFLOPS", flops as f64 / 1_000.0)
551            } else {
552                format!("{} FLOPS", flops)
553            }
554        }
555    }
556
557    /// Advanced model analyzer with detailed metrics
558    #[derive(Debug, Clone)]
559    pub struct ModelAnalyzer {
560        pub config: AnalysisConfig,
561    }
562
563    #[derive(Debug, Clone)]
564    pub struct AnalysisConfig {
565        pub analyze_gradients: bool,
566        pub analyze_activations: bool,
567        pub analyze_flops: bool,
568        pub analyze_memory: bool,
569        pub batch_analysis: bool,
570    }
571
572    impl Default for AnalysisConfig {
573        fn default() -> Self {
574            Self {
575                analyze_gradients: false,
576                analyze_activations: false,
577                analyze_flops: true,
578                analyze_memory: true,
579                batch_analysis: false,
580            }
581        }
582    }
583
584    impl ModelAnalyzer {
585        pub fn new(config: AnalysisConfig) -> Self {
586            Self { config }
587        }
588
589        pub fn default() -> Self {
590            Self::new(AnalysisConfig::default())
591        }
592
593        /// Comprehensive model analysis
594        pub fn analyze<M: Module>(
595            &self,
596            model: &M,
597            input_shape: &[usize],
598        ) -> Result<AnalysisReport> {
599            let mut report = AnalysisReport::new();
600
601            // Basic model information
602            let parameters = model.parameters();
603            report.parameter_count = count_parameters(&parameters);
604            report.model_size_mb = (report.parameter_count * 4) as f64 / 1_048_576.0;
605
606            // Memory analysis
607            if self.config.analyze_memory {
608                report.memory_analysis = Some(self.analyze_memory(model, input_shape)?);
609            }
610
611            // FLOPS analysis
612            if self.config.analyze_flops {
613                report.flops_analysis = Some(self.estimate_flops(model, input_shape)?);
614            }
615
616            Ok(report)
617        }
618
619        fn analyze_memory<M: Module>(
620            &self,
621            model: &M,
622            input_shape: &[usize],
623        ) -> Result<MemoryAnalysis> {
624            let input_memory = estimate_tensor_memory(input_shape, 4);
625            let param_memory = count_parameters(&model.parameters()) * 4;
626
627            // Estimate intermediate activations (rough approximation)
628            let intermediate_memory = input_memory * 3; // Rough estimate
629
630            Ok(MemoryAnalysis {
631                input_memory,
632                parameter_memory: param_memory,
633                intermediate_memory,
634                total_memory: input_memory + param_memory + intermediate_memory,
635            })
636        }
637
638        fn estimate_flops<M: Module>(
639            &self,
640            _model: &M,
641            input_shape: &[usize],
642        ) -> Result<FLOPSAnalysis> {
643            // Simplified FLOPS estimation
644            // In a full implementation, this would traverse the model structure
645            let estimated_flops = input_shape.iter().product::<usize>() as u64 * 1000; // Rough estimate
646
647            Ok(FLOPSAnalysis {
648                total_flops: estimated_flops,
649                flops_per_layer: HashMap::new(),
650            })
651        }
652    }
653
654    /// Memory analysis results
655    #[derive(Debug, Clone)]
656    pub struct MemoryAnalysis {
657        pub input_memory: usize,
658        pub parameter_memory: usize,
659        pub intermediate_memory: usize,
660        pub total_memory: usize,
661    }
662
663    /// FLOPS analysis results
664    #[derive(Debug, Clone)]
665    pub struct FLOPSAnalysis {
666        pub total_flops: u64,
667        pub flops_per_layer: HashMap<String, u64>,
668    }
669
670    /// Comprehensive analysis report
671    #[derive(Debug, Clone)]
672    pub struct AnalysisReport {
673        pub parameter_count: usize,
674        pub model_size_mb: f64,
675        pub memory_analysis: Option<MemoryAnalysis>,
676        pub flops_analysis: Option<FLOPSAnalysis>,
677    }
678
679    impl AnalysisReport {
680        pub fn new() -> Self {
681            Self {
682                parameter_count: 0,
683                model_size_mb: 0.0,
684                memory_analysis: None,
685                flops_analysis: None,
686            }
687        }
688
689        /// Print detailed analysis report
690        pub fn print_detailed(&self) {
691            println!("=== Detailed Model Analysis ===");
692            println!("Parameters: {}", format_number(self.parameter_count));
693            println!("Model Size: {:.2} MB", self.model_size_mb);
694
695            if let Some(memory) = &self.memory_analysis {
696                println!("\n--- Memory Analysis ---");
697                println!("Input Memory: {}", format_bytes(memory.input_memory));
698                println!(
699                    "Parameter Memory: {}",
700                    format_bytes(memory.parameter_memory)
701                );
702                println!(
703                    "Intermediate Memory: {}",
704                    format_bytes(memory.intermediate_memory)
705                );
706                println!("Total Memory: {}", format_bytes(memory.total_memory));
707            }
708
709            if let Some(flops) = &self.flops_analysis {
710                println!("\n--- FLOPS Analysis ---");
711                println!(
712                    "Total FLOPS: {}",
713                    FLOPSCounter::format_flops(flops.total_flops)
714                );
715
716                if !flops.flops_per_layer.is_empty() {
717                    println!("Per-layer FLOPS:");
718                    for (layer, flops) in &flops.flops_per_layer {
719                        println!("  {}: {}", layer, FLOPSCounter::format_flops(*flops));
720                    }
721                }
722            }
723        }
724    }
725
726    /// Batch profiler for statistical analysis
727    pub struct BatchProfiler {
728        config: BatchProfilingConfig,
729    }
730
731    #[derive(Debug, Clone)]
732    pub struct BatchProfilingConfig {
733        pub num_runs: usize,
734        pub warmup_runs: usize,
735        pub collect_stats: bool,
736    }
737
738    impl Default for BatchProfilingConfig {
739        fn default() -> Self {
740            Self {
741                num_runs: 10,
742                warmup_runs: 3,
743                collect_stats: true,
744            }
745        }
746    }
747
748    impl BatchProfiler {
749        pub fn new(config: BatchProfilingConfig) -> Self {
750            Self { config }
751        }
752
753        /// Run batch profiling on a model
754        #[cfg(feature = "std")]
755        pub fn profile_batch<M: Module>(
756            &self,
757            model: &M,
758            input: &Tensor,
759        ) -> Result<BatchProfilingResult> {
760            let mut times = Vec::new();
761
762            // Warmup runs
763            for _ in 0..self.config.warmup_runs {
764                let _output = model.forward(input)?;
765            }
766
767            // Actual profiling runs
768            for _ in 0..self.config.num_runs {
769                let start = Instant::now();
770                let _output = model.forward(input)?;
771                let elapsed = start.elapsed();
772                times.push(elapsed);
773            }
774
775            Ok(BatchProfilingResult::from_times(times))
776        }
777    }
778
779    /// Results from batch profiling
780    #[derive(Debug, Clone)]
781    pub struct BatchProfilingResult {
782        pub mean_time: f64,
783        pub std_time: f64,
784        pub min_time: f64,
785        pub max_time: f64,
786        pub median_time: f64,
787        pub num_runs: usize,
788    }
789
790    #[cfg(feature = "std")]
791    impl BatchProfilingResult {
792        pub fn from_times(times: Vec<Duration>) -> Self {
793            let times_ms: Vec<f64> = times.iter().map(|d| d.as_secs_f64() * 1000.0).collect();
794
795            let mean = times_ms.iter().sum::<f64>() / times_ms.len() as f64;
796            let variance =
797                times_ms.iter().map(|t| (t - mean).powi(2)).sum::<f64>() / times_ms.len() as f64;
798            let std_dev = variance.sqrt();
799
800            let mut sorted_times = times_ms.clone();
801            sorted_times
802                .sort_by(|a, b| a.partial_cmp(b).expect("comparison should not involve NaN"));
803
804            let median = if sorted_times.len() % 2 == 0 {
805                (sorted_times[sorted_times.len() / 2 - 1] + sorted_times[sorted_times.len() / 2])
806                    / 2.0
807            } else {
808                sorted_times[sorted_times.len() / 2]
809            };
810
811            Self {
812                mean_time: mean,
813                std_time: std_dev,
814                min_time: sorted_times[0],
815                max_time: sorted_times[sorted_times.len() - 1],
816                median_time: median,
817                num_runs: times.len(),
818            }
819        }
820
821        pub fn print_stats(&self) {
822            println!("=== Batch Profiling Results ===");
823            println!("Runs: {}", self.num_runs);
824            println!("Mean: {:.3}ms", self.mean_time);
825            println!("Std Dev: {:.3}ms", self.std_time);
826            println!("Min: {:.3}ms", self.min_time);
827            println!("Max: {:.3}ms", self.max_time);
828            println!("Median: {:.3}ms", self.median_time);
829        }
830    }
831}
832
833#[cfg(test)]
834mod tests {
835    use super::*;
836    use crate::layers::Linear;
837    use torsh_tensor::creation::randn;
838
839    #[test]
840    fn test_layer_info_creation() {
841        let layer_info = LayerInfo::new(
842            "linear1".to_string(),
843            "Linear".to_string(),
844            vec![10, 20],
845            vec![10, 30],
846            600, // 20 * 30 weights + 30 biases
847            600,
848        );
849
850        assert_eq!(layer_info.name, "linear1");
851        assert_eq!(layer_info.layer_type, "Linear");
852        assert_eq!(layer_info.param_count, 600);
853        assert_eq!(layer_info.trainable_params, 600);
854    }
855
856    #[test]
857    fn test_format_number() {
858        assert_eq!(format_number(500), "500");
859        assert_eq!(format_number(1500), "1.5K");
860        assert_eq!(format_number(1_500_000), "1.5M");
861        assert_eq!(format_number(1_500_000_000), "1.5B");
862    }
863
864    #[test]
865    fn test_format_bytes() {
866        assert_eq!(format_bytes(500), "500 B");
867        assert_eq!(format_bytes(1536), "1.5 KB");
868        assert_eq!(format_bytes(1_572_864), "1.5 MB");
869        assert_eq!(format_bytes(1_610_612_736), "1.5 GB");
870    }
871
872    #[test]
873    fn test_model_summary() -> Result<()> {
874        let model = Linear::new(128, 64, true);
875        let input_shape = [10, 128];
876
877        let summary = summarize(&model, &input_shape, None)?;
878
879        assert_eq!(summary.input_shape, vec![10, 128]);
880        assert!(!summary.layers.is_empty());
881        assert!(summary.total_params > 0);
882
883        Ok(())
884    }
885
886    #[test]
887    fn test_model_profiler() -> Result<()> {
888        let model = Linear::new(64, 32, true);
889        let input = randn::<f32>(&[8, 64])?;
890
891        let profiler = ModelProfiler::new().with_time_tracking();
892        let result = profiler.profile(&model, &input)?;
893
894        assert_eq!(result.input_shape, vec![8, 64]);
895        assert_eq!(result.output_shape, vec![8, 32]);
896        assert!(result.parameter_count > 0);
897
898        Ok(())
899    }
900
901    #[test]
902    fn test_utils_functions() -> Result<()> {
903        let model = Linear::new(100, 50, true);
904
905        let param_count = utils::count_model_parameters(&model);
906        assert_eq!(param_count, 100 * 50 + 50); // weights + biases
907
908        let model_size = utils::get_model_size_mb(&model);
909        assert!(model_size > 0.0);
910
911        let fits_budget = utils::check_memory_budget(&model, &[10, 100], 100.0);
912        assert!(fits_budget); // Should easily fit in 100MB
913
914        Ok(())
915    }
916
917    #[test]
918    fn test_flops_counter() {
919        let mut counter = profiling::FLOPSCounter::new();
920
921        // Test linear layer FLOPS counting
922        counter.count_linear_flops("linear1".to_string(), 128, 64, 32);
923        assert_eq!(counter.total_flops, 32 * 128 * 64 * 2); // batch * input * output * 2
924
925        // Test FLOPS formatting
926        assert_eq!(profiling::FLOPSCounter::format_flops(1500), "1.50 KFLOPS");
927        assert_eq!(
928            profiling::FLOPSCounter::format_flops(1_500_000),
929            "1.50 MFLOPS"
930        );
931        assert_eq!(
932            profiling::FLOPSCounter::format_flops(1_500_000_000),
933            "1.50 GFLOPS"
934        );
935    }
936
937    #[test]
938    fn test_model_analyzer() -> Result<()> {
939        let model = Linear::new(128, 64, true);
940        let input_shape = [10, 128];
941
942        let analyzer = profiling::ModelAnalyzer::default();
943        let report = analyzer.analyze(&model, &input_shape)?;
944
945        assert!(report.parameter_count > 0);
946        assert!(report.model_size_mb > 0.0);
947        assert!(report.memory_analysis.is_some());
948        assert!(report.flops_analysis.is_some());
949
950        if let Some(memory) = &report.memory_analysis {
951            assert!(memory.total_memory > 0);
952            assert!(memory.parameter_memory > 0);
953        }
954
955        Ok(())
956    }
957
958    #[test]
959    fn test_analysis_config() {
960        let config = profiling::AnalysisConfig::default();
961        assert!(!config.analyze_gradients);
962        assert!(!config.analyze_activations);
963        assert!(config.analyze_flops);
964        assert!(config.analyze_memory);
965        assert!(!config.batch_analysis);
966    }
967
968    #[test]
969    fn test_batch_profiling_config() {
970        let config = profiling::BatchProfilingConfig::default();
971        assert_eq!(config.num_runs, 10);
972        assert_eq!(config.warmup_runs, 3);
973        assert!(config.collect_stats);
974    }
975
976    #[cfg(feature = "std")]
977    #[test]
978    fn test_batch_profiler() -> Result<()> {
979        let model = Linear::new(64, 32, true);
980        let input = randn::<f32>(&[8, 64])?;
981
982        let config = profiling::BatchProfilingConfig {
983            num_runs: 5,
984            warmup_runs: 2,
985            collect_stats: true,
986        };
987
988        let profiler = profiling::BatchProfiler::new(config);
989        let result = profiler.profile_batch(&model, &input)?;
990
991        assert_eq!(result.num_runs, 5);
992        assert!(result.mean_time >= 0.0);
993        assert!(result.std_time >= 0.0);
994        assert!(result.min_time >= 0.0);
995        assert!(result.max_time >= result.min_time);
996        assert!(result.median_time >= 0.0);
997
998        Ok(())
999    }
1000}
torsh_nn/summary.rs

torsh_nn/
summary.rs
