torsh-cli 0.1.2

Command-line tools for the ToRSh deep learning framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105

//! Model operation commands
//!
//! This module provides a comprehensive set of model operations including:
//! - Conversion between different model formats
//! - Optimization for deployment
//! - Quantization for efficiency
//! - Pruning for size reduction
//! - Analysis and inspection
//! - Benchmarking performance
//!
//! The implementation is modularized for better maintainability:
//! - `analysis`: Model inspection and validation
//! - `benchmarking`: Performance testing
//! - `conversion`: Format conversion and utility operations
//! - `optimization`: Model optimization, quantization, and pruning
//! - `types`: Shared data structures
//! - `args`: Command-line argument definitions

use anyhow::Result;
use clap::Subcommand;

use crate::config::Config;

// Import all sub-modules
pub mod analysis;
pub mod args;
pub mod benchmarking;
pub mod conversion;
pub mod enhanced_profiling;
pub mod enhanced_serialization;
pub mod optimization;
pub mod profiling;
pub mod pytorch_parser;
pub mod real_benchmarking;
pub mod serialization;
pub mod tensor_integration;
pub mod types;
pub mod validation;

// Re-export commonly used types for convenience
pub use args::*;

/// Model operation subcommands
#[derive(Subcommand)]
pub enum ModelCommands {
    /// Convert model between different formats
    Convert(ConvertArgs),

    /// Optimize model for deployment
    Optimize(OptimizeArgs),

    /// Quantize model to reduce size and improve performance
    Quantize(QuantizeArgs),

    /// Prune model to remove unnecessary parameters
    Prune(PruneArgs),

    /// Inspect model architecture and properties
    Inspect(InspectArgs),

    /// Validate model functionality and accuracy
    Validate(ValidateArgs),

    /// Benchmark model performance
    Benchmark(BenchmarkArgs),

    /// Compress model using various techniques
    Compress(CompressArgs),

    /// Extract model components (weights, architecture, etc.)
    Extract(ExtractArgs),

    /// Merge multiple models
    Merge(MergeArgs),
}

/// Execute model operation based on subcommand
pub async fn execute(cmd: ModelCommands, config: &Config, output_format: &str) -> Result<()> {
    match cmd {
        ModelCommands::Convert(args) => {
            conversion::convert_model(args, config, output_format).await
        }
        ModelCommands::Optimize(args) => {
            optimization::optimize_model(args, config, output_format).await
        }
        ModelCommands::Quantize(args) => {
            optimization::quantize_model(args, config, output_format).await
        }
        ModelCommands::Prune(args) => optimization::prune_model(args, config, output_format).await,
        ModelCommands::Inspect(args) => analysis::inspect_model(args, config, output_format).await,
        ModelCommands::Validate(args) => {
            analysis::validate_model(args, config, output_format).await
        }
        ModelCommands::Benchmark(args) => {
            benchmarking::benchmark_model(args, config, output_format).await
        }
        ModelCommands::Compress(args) => {
            conversion::compress_model(args, config, output_format).await
        }
        ModelCommands::Extract(args) => {
            conversion::extract_model(args, config, output_format).await
        }
        ModelCommands::Merge(args) => conversion::merge_model(args, config, output_format).await,
    }
}