Files
rustytorch/crates/tooling/rtx-bench/src/models/vision.rs
T
2026-03-04 00:08:42 +00:00

85 lines
2.2 KiB
Rust

//! Vision model benchmarks
use crate::{BenchmarkConfig, BenchmarkMeasurement, time_benchmark_async};
use anyhow::Result;
use std::collections::HashMap;
pub struct VisionBenchmarks;
impl VisionBenchmarks {
pub fn new() -> Self {
Self
}
pub async fn run_all_benchmarks(
&self,
config: &BenchmarkConfig,
) -> Result<HashMap<String, Vec<BenchmarkMeasurement>>> {
let mut measurements = HashMap::new();
measurements.insert(
"resnet50_inference".to_string(),
self.benchmark_resnet_inference(config).await?,
);
measurements.insert(
"vit_inference".to_string(),
self.benchmark_vit_inference(config).await?,
);
Ok(measurements)
}
async fn benchmark_resnet_inference(
&self,
config: &BenchmarkConfig,
) -> Result<Vec<BenchmarkMeasurement>> {
let mut measurements = Vec::with_capacity(config.measurement_iterations);
for _ in 0..config.measurement_iterations {
let measurement = time_benchmark_async("resnet50_inference".to_string(), || async {
self.resnet_forward().await?;
Ok(())
})
.await?;
measurements.push(measurement);
}
Ok(measurements)
}
async fn benchmark_vit_inference(
&self,
config: &BenchmarkConfig,
) -> Result<Vec<BenchmarkMeasurement>> {
let mut measurements = Vec::with_capacity(config.measurement_iterations);
for _ in 0..config.measurement_iterations {
let measurement = time_benchmark_async("vit_inference".to_string(), || async {
self.vit_forward().await?;
Ok(())
})
.await?;
measurements.push(measurement);
}
Ok(measurements)
}
async fn resnet_forward(&self) -> Result<()> {
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
Ok(())
}
async fn vit_forward(&self) -> Result<()> {
tokio::time::sleep(std::time::Duration::from_millis(80)).await;
Ok(())
}
}
impl Default for VisionBenchmarks {
fn default() -> Self {
Self::new()
}
}