46 lines
1.6 KiB
Rust
46 lines
1.6 KiB
Rust
//! Criterion benchmarks for model architectures
|
|
|
|
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
|
|
use rtx_bench::{BenchmarkConfig, models::transformers::TransformerBenchmarks};
|
|
use tokio::runtime::Runtime;
|
|
|
|
fn benchmark_transformer_models(c: &mut Criterion) {
|
|
let rt = Runtime::new().unwrap();
|
|
let benchmarks = TransformerBenchmarks::new();
|
|
let config = BenchmarkConfig::default()
|
|
.with_measurement_iterations(5)
|
|
.with_warmup_iterations(2);
|
|
|
|
let mut group = c.benchmark_group("transformer_models");
|
|
group.sample_size(10);
|
|
|
|
// Test with smaller configurations for fast benchmarking
|
|
for batch_size in [1, 4, 8].iter() {
|
|
group.bench_with_input(
|
|
BenchmarkId::new("bert_forward", batch_size),
|
|
batch_size,
|
|
|b, &batch_size| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let model_config = rtx_bench::models::transformers::ModelConfig {
|
|
model_type: rtx_bench::models::transformers::ModelType::Bert,
|
|
hidden_size: 256,
|
|
num_layers: 4,
|
|
num_heads: 8,
|
|
sequence_length: 128,
|
|
vocab_size: 10000,
|
|
};
|
|
benchmarks
|
|
.benchmark_forward_pass(&config, &model_config, batch_size)
|
|
.await
|
|
.unwrap();
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(benches, benchmark_transformer_models);
|
|
criterion_main!(benches);
|