449 lines
12 KiB
Rust
449 lines
12 KiB
Rust
//! Configuration types for inference profiling.
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Model architecture type for profiling.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub enum ModelType {
|
|
/// ResNet-18 (11M parameters)
|
|
#[serde(rename = "resnet18")]
|
|
ResNet18,
|
|
/// ResNet-50 (25M parameters)
|
|
#[serde(rename = "resnet50")]
|
|
ResNet50,
|
|
/// Vision Transformer Base/16 (86M parameters)
|
|
#[serde(rename = "vit-b16")]
|
|
ViTB16,
|
|
/// Vision Transformer Large/16 (304M parameters)
|
|
#[serde(rename = "vit-l16")]
|
|
ViTL16,
|
|
/// `ConvNeXt` Tiny (29M parameters)
|
|
#[serde(rename = "convnext-tiny")]
|
|
ConvNeXtTiny,
|
|
/// `ConvNeXt` Base (89M parameters)
|
|
#[serde(rename = "convnext-base")]
|
|
ConvNeXtBase,
|
|
/// Custom user-provided model
|
|
#[serde(rename = "custom")]
|
|
Custom,
|
|
}
|
|
|
|
/// Compute device type.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
#[serde(rename_all = "UPPERCASE")]
|
|
pub enum DeviceType {
|
|
/// CPU execution
|
|
CPU,
|
|
/// CUDA GPU execution
|
|
CUDA,
|
|
/// Metal GPU execution (macOS)
|
|
Metal,
|
|
}
|
|
|
|
/// Input shape for profiling (NCHW format).
|
|
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
|
pub struct InputShape {
|
|
/// Number of channels
|
|
pub channels: usize,
|
|
/// Height
|
|
pub height: usize,
|
|
/// Width
|
|
pub width: usize,
|
|
}
|
|
|
|
/// Profile configuration.
|
|
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
|
pub struct ProfileConfig {
|
|
/// Model to profile
|
|
pub model_type: ModelType,
|
|
/// Device to run on
|
|
pub device: DeviceType,
|
|
/// Batch sizes to test
|
|
pub batch_sizes: Vec<usize>,
|
|
/// Warmup iterations before measurement
|
|
pub warmup_iterations: usize,
|
|
/// Benchmark iterations for measurement
|
|
pub benchmark_iterations: usize,
|
|
/// Input shape (excluding batch dimension)
|
|
pub input_shape: InputShape,
|
|
}
|
|
|
|
impl ProfileConfig {
|
|
/// Creates a new profile configuration.
|
|
///
|
|
/// # Errors
|
|
/// Returns an error if validation fails.
|
|
pub fn new(
|
|
model_type: ModelType,
|
|
device: DeviceType,
|
|
batch_sizes: Vec<usize>,
|
|
warmup_iterations: usize,
|
|
benchmark_iterations: usize,
|
|
input_shape: InputShape,
|
|
) -> Result<Self, String> {
|
|
let config = Self {
|
|
model_type,
|
|
device,
|
|
batch_sizes,
|
|
warmup_iterations,
|
|
benchmark_iterations,
|
|
input_shape,
|
|
};
|
|
config.validate()?;
|
|
Ok(config)
|
|
}
|
|
|
|
/// Validates the configuration.
|
|
///
|
|
/// # Errors
|
|
/// Returns an error if any validation constraint fails.
|
|
pub fn validate(&self) -> Result<(), String> {
|
|
if self.batch_sizes.is_empty() {
|
|
return Err("batch_sizes cannot be empty".to_string());
|
|
}
|
|
|
|
for &batch_size in &self.batch_sizes {
|
|
if batch_size == 0 {
|
|
return Err("batch_size cannot be zero".to_string());
|
|
}
|
|
if batch_size > 1024 {
|
|
return Err(format!("batch_size {batch_size} exceeds maximum of 1024"));
|
|
}
|
|
}
|
|
|
|
if self.warmup_iterations > 10000 {
|
|
return Err(format!(
|
|
"warmup_iterations {} exceeds maximum of 10000",
|
|
self.warmup_iterations
|
|
));
|
|
}
|
|
|
|
if self.benchmark_iterations == 0 {
|
|
return Err("benchmark_iterations cannot be zero".to_string());
|
|
}
|
|
|
|
if self.benchmark_iterations > 100000 {
|
|
return Err(format!(
|
|
"benchmark_iterations {} exceeds maximum of 100000",
|
|
self.benchmark_iterations
|
|
));
|
|
}
|
|
|
|
if self.input_shape.channels == 0 {
|
|
return Err("input_shape.channels cannot be zero".to_string());
|
|
}
|
|
|
|
if self.input_shape.height == 0 {
|
|
return Err("input_shape.height cannot be zero".to_string());
|
|
}
|
|
|
|
if self.input_shape.width == 0 {
|
|
return Err("input_shape.width cannot be zero".to_string());
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Default for InputShape {
|
|
fn default() -> Self {
|
|
Self {
|
|
channels: 3,
|
|
height: 224,
|
|
width: 224,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for ProfileConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
model_type: ModelType::ResNet18,
|
|
device: DeviceType::CPU,
|
|
batch_sizes: vec![1, 2, 4, 8],
|
|
warmup_iterations: 10,
|
|
benchmark_iterations: 100,
|
|
input_shape: InputShape::default(),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_model_type_serialization() {
|
|
let model = ModelType::ResNet18;
|
|
let json = serde_json::to_string(&model).expect("serialization failed");
|
|
assert_eq!(json, "\"resnet18\"");
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_type_deserialization() {
|
|
let json = "\"vit-b16\"";
|
|
let model: ModelType = serde_json::from_str(json).expect("deserialization failed");
|
|
assert_eq!(model, ModelType::ViTB16);
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_type_roundtrip() {
|
|
let models = vec![
|
|
ModelType::ResNet18,
|
|
ModelType::ResNet50,
|
|
ModelType::ViTB16,
|
|
ModelType::ViTL16,
|
|
ModelType::ConvNeXtTiny,
|
|
ModelType::ConvNeXtBase,
|
|
ModelType::Custom,
|
|
];
|
|
|
|
for model in models {
|
|
let json = serde_json::to_string(&model).expect("serialization failed");
|
|
let decoded: ModelType = serde_json::from_str(&json).expect("deserialization failed");
|
|
assert_eq!(model, decoded);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_device_type_serialization() {
|
|
let device = DeviceType::CUDA;
|
|
let json = serde_json::to_string(&device).expect("serialization failed");
|
|
assert_eq!(json, "\"CUDA\"");
|
|
}
|
|
|
|
#[test]
|
|
fn test_device_type_deserialization() {
|
|
let json = "\"CPU\"";
|
|
let device: DeviceType = serde_json::from_str(json).expect("deserialization failed");
|
|
assert_eq!(device, DeviceType::CPU);
|
|
}
|
|
|
|
#[test]
|
|
fn test_device_type_roundtrip() {
|
|
let devices = vec![DeviceType::CPU, DeviceType::CUDA, DeviceType::Metal];
|
|
|
|
for device in devices {
|
|
let json = serde_json::to_string(&device).expect("serialization failed");
|
|
let decoded: DeviceType = serde_json::from_str(&json).expect("deserialization failed");
|
|
assert_eq!(device, decoded);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_success() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2, 4],
|
|
10,
|
|
100,
|
|
InputShape {
|
|
channels: 3,
|
|
height: 224,
|
|
width: 224,
|
|
},
|
|
);
|
|
|
|
assert!(config.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_empty_batch_sizes() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![],
|
|
10,
|
|
100,
|
|
InputShape::default(),
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
assert_eq!(config.unwrap_err(), "batch_sizes cannot be empty");
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_zero_batch_size() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 0, 4],
|
|
10,
|
|
100,
|
|
InputShape::default(),
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
assert_eq!(config.unwrap_err(), "batch_size cannot be zero");
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_batch_size_too_large() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2048],
|
|
10,
|
|
100,
|
|
InputShape::default(),
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
assert!(config.unwrap_err().contains("exceeds maximum of 1024"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_warmup_iterations_too_large() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2],
|
|
20000,
|
|
100,
|
|
InputShape::default(),
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
let err = config.unwrap_err();
|
|
assert!(err.contains("warmup_iterations"));
|
|
assert!(err.contains("exceeds maximum of 10000"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_zero_benchmark_iterations() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2],
|
|
10,
|
|
0,
|
|
InputShape::default(),
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
assert_eq!(config.unwrap_err(), "benchmark_iterations cannot be zero");
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_benchmark_iterations_too_large() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2],
|
|
10,
|
|
200000,
|
|
InputShape::default(),
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
let err = config.unwrap_err();
|
|
assert!(err.contains("benchmark_iterations"));
|
|
assert!(err.contains("exceeds maximum of 100000"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_zero_channels() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2],
|
|
10,
|
|
100,
|
|
InputShape {
|
|
channels: 0,
|
|
height: 224,
|
|
width: 224,
|
|
},
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
assert_eq!(config.unwrap_err(), "input_shape.channels cannot be zero");
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_zero_height() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2],
|
|
10,
|
|
100,
|
|
InputShape {
|
|
channels: 3,
|
|
height: 0,
|
|
width: 224,
|
|
},
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
assert_eq!(config.unwrap_err(), "input_shape.height cannot be zero");
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_validation_zero_width() {
|
|
let config = ProfileConfig::new(
|
|
ModelType::ResNet18,
|
|
DeviceType::CPU,
|
|
vec![1, 2],
|
|
10,
|
|
100,
|
|
InputShape {
|
|
channels: 3,
|
|
height: 224,
|
|
width: 0,
|
|
},
|
|
);
|
|
|
|
assert!(config.is_err());
|
|
assert_eq!(config.unwrap_err(), "input_shape.width cannot be zero");
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_serialization() {
|
|
let config = ProfileConfig::default();
|
|
let json = serde_json::to_string(&config).expect("serialization failed");
|
|
assert!(json.contains("\"model_type\""));
|
|
assert!(json.contains("\"resnet18\""));
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_deserialization() {
|
|
let json = r#"{
|
|
"model_type": "vit-b16",
|
|
"device": "CUDA",
|
|
"batch_sizes": [1, 2, 4],
|
|
"warmup_iterations": 10,
|
|
"benchmark_iterations": 100,
|
|
"input_shape": {
|
|
"channels": 3,
|
|
"height": 224,
|
|
"width": 224
|
|
}
|
|
}"#;
|
|
|
|
let config: ProfileConfig = serde_json::from_str(json).expect("deserialization failed");
|
|
assert_eq!(config.model_type, ModelType::ViTB16);
|
|
assert_eq!(config.device, DeviceType::CUDA);
|
|
assert_eq!(config.batch_sizes, vec![1, 2, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_input_shape_default() {
|
|
let shape = InputShape::default();
|
|
assert_eq!(shape.channels, 3);
|
|
assert_eq!(shape.height, 224);
|
|
assert_eq!(shape.width, 224);
|
|
}
|
|
|
|
#[test]
|
|
fn test_profile_config_default() {
|
|
let config = ProfileConfig::default();
|
|
assert_eq!(config.model_type, ModelType::ResNet18);
|
|
assert_eq!(config.device, DeviceType::CPU);
|
|
assert_eq!(config.batch_sizes, vec![1, 2, 4, 8]);
|
|
assert_eq!(config.warmup_iterations, 10);
|
|
assert_eq!(config.benchmark_iterations, 100);
|
|
}
|
|
}
|