108 lines
3.1 KiB
Rust
108 lines
3.1 KiB
Rust
use criterion::{Criterion, black_box, criterion_group, criterion_main};
|
|
use rtx_multimodal::*;
|
|
use rtx_runtime::Device;
|
|
use rtx_tensor::Tensor;
|
|
|
|
fn benchmark_vit_forward(c: &mut Criterion) {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = vision::vit::ViTConfig {
|
|
image_size: 224,
|
|
patch_size: 16,
|
|
in_channels: 3,
|
|
embed_dim: 768,
|
|
depth: 12,
|
|
num_heads: 12,
|
|
mlp_ratio: 4.0,
|
|
num_classes: 1000,
|
|
dropout: 0.0,
|
|
attention_dropout: 0.0,
|
|
};
|
|
|
|
let model = vision::vit::VisionTransformer::new(&config, &device).unwrap();
|
|
let input = Tensor::randn(&[1, 3, 224, 224], &device).unwrap();
|
|
|
|
c.bench_function("vit_forward", |b| {
|
|
b.iter(|| {
|
|
let _output = model.forward(black_box(&input)).unwrap();
|
|
})
|
|
});
|
|
}
|
|
|
|
fn benchmark_conformer_forward(c: &mut Criterion) {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = audio::conformer::ConformerEncoderConfig {
|
|
input_dim: 80,
|
|
d_model: 512,
|
|
num_layers: 6,
|
|
num_heads: 8,
|
|
feed_forward_expansion_factor: 4,
|
|
conv_kernel_size: 31,
|
|
dropout: 0.0,
|
|
};
|
|
|
|
let model = audio::conformer::ConformerEncoder::new(&config, &device).unwrap();
|
|
let input = Tensor::randn(&[1, 100, 80], &device).unwrap();
|
|
|
|
c.bench_function("conformer_forward", |b| {
|
|
b.iter(|| {
|
|
let _output = model.forward(black_box(&input)).unwrap();
|
|
})
|
|
});
|
|
}
|
|
|
|
fn benchmark_timesformer_forward(c: &mut Criterion) {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = video::timesformer::TimeSformerConfig {
|
|
image_size: 224,
|
|
patch_size: 16,
|
|
in_channels: 3,
|
|
embed_dim: 768,
|
|
depth: 12,
|
|
num_heads: 12,
|
|
mlp_ratio: 4.0,
|
|
num_frames: 8,
|
|
num_classes: 400,
|
|
dropout: 0.0,
|
|
attention_dropout: 0.0,
|
|
};
|
|
|
|
let model = video::timesformer::TimeSformerModel::new(&config, &device).unwrap();
|
|
let input = Tensor::randn(&[1, 3, 8, 224, 224], &device).unwrap();
|
|
|
|
c.bench_function("timesformer_forward", |b| {
|
|
b.iter(|| {
|
|
let _output = model.forward(black_box(&input)).unwrap();
|
|
})
|
|
});
|
|
}
|
|
|
|
fn benchmark_cross_modal_attention(c: &mut Criterion) {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = fusion::CrossModalAttentionConfig {
|
|
embed_dim: 512,
|
|
num_heads: 8,
|
|
dropout: 0.0,
|
|
};
|
|
|
|
let attention = fusion::CrossModalAttention::new(&config, &device).unwrap();
|
|
let vision = Tensor::randn(&[1, 197, 512], &device).unwrap();
|
|
let audio = Tensor::randn(&[1, 100, 512], &device).unwrap();
|
|
|
|
c.bench_function("cross_modal_attention", |b| {
|
|
b.iter(|| {
|
|
let _output = attention
|
|
.forward(black_box(&vision), black_box(&audio))
|
|
.unwrap();
|
|
})
|
|
});
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
benchmark_vit_forward,
|
|
benchmark_conformer_forward,
|
|
benchmark_timesformer_forward,
|
|
benchmark_cross_modal_attention
|
|
);
|
|
criterion_main!(benches);
|