Files
rustytorch/crates/models/rtx-vision-advanced/benches/segmentation_bench.rs
T
2026-03-04 00:08:42 +00:00

422 lines
14 KiB
Rust

//! Segmentation model benchmarks
//!
//! Performance benchmarks for:
//! - Semantic segmentation (DeepLabV3+)
//! - Instance segmentation (Mask R-CNN)
//! - Panoptic segmentation
//! - Medical image segmentation
//! - Real-time segmentation optimization
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rtx_tensor::{DType, Device, Tensor};
use rtx_vision_advanced::*;
/// Benchmark helper for segmentation tests
struct SegmentationBenchHelper;
impl SegmentationBenchHelper {
fn create_image(batch_size: usize, height: usize, width: usize) -> Tensor {
Tensor::randn(&[batch_size, 3, height, width], &Device::default())
.expect("Failed to create benchmark image")
}
fn create_medical_volume(depth: usize, height: usize, width: usize) -> Tensor {
Tensor::randn(&[1, depth, height, width], &Device::default())
.expect("Failed to create benchmark medical volume")
}
}
/// Benchmark semantic segmentation models
fn bench_semantic_segmentation(c: &mut Criterion) {
let mut group = c.benchmark_group("semantic_segmentation");
// DeepLabV3+ with different backbones
let backbones = vec![("resnet50", "resnet50"), ("resnet101", "resnet101")];
for (name, backbone) in backbones {
let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, backbone)
.expect("Failed to create DeepLabV3+");
let test_image = SegmentationBenchHelper::create_image(1, 512, 512);
let config = segmentation::SegmentationConfig {
segmentation_type: segmentation::SegmentationType::Semantic,
..Default::default()
};
group.throughput(Throughput::Elements(512 * 512)); // Pixels
group.bench_with_input(
BenchmarkId::new("deeplabv3plus", name),
&(test_image, config),
|b, (image, cfg)| {
b.iter(|| segmentor.segment(image, cfg).expect("Segmentation failed"));
},
);
}
group.finish();
}
/// Benchmark instance segmentation
fn bench_instance_segmentation(c: &mut Criterion) {
let mut group = c.benchmark_group("instance_segmentation");
let mut mask_rcnn = segmentation::SegmentationFactory::create_mask_rcnn(80)
.expect("Failed to create Mask R-CNN");
let input_sizes = vec![
("small", 512, 512),
("medium", 800, 1333),
("large", 1024, 1024),
];
for (size_name, height, width) in input_sizes {
let test_image = SegmentationBenchHelper::create_image(1, height, width);
let config = segmentation::SegmentationConfig {
segmentation_type: segmentation::SegmentationType::Instance,
vision_config: VisionConfig {
input_size: (height, width),
..Default::default()
},
..Default::default()
};
group.throughput(Throughput::Elements((height * width) as u64));
group.bench_with_input(
BenchmarkId::new("mask_rcnn", size_name),
&(test_image, config),
|b, (image, cfg)| {
b.iter(|| {
mask_rcnn
.segment(image, cfg)
.expect("Instance segmentation failed")
});
},
);
}
group.finish();
}
/// Benchmark multi-scale segmentation
fn bench_multiscale_segmentation(c: &mut Criterion) {
let mut group = c.benchmark_group("multiscale_segmentation");
let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
.expect("Failed to create segmentor");
let scale_configs = vec![
("single_scale", vec![1.0]),
("multi_scale_3", vec![0.75, 1.0, 1.25]),
("multi_scale_5", vec![0.5, 0.75, 1.0, 1.25, 1.5]),
];
let test_image = SegmentationBenchHelper::create_image(1, 512, 512);
for (config_name, scales) in scale_configs {
let config = segmentation::SegmentationConfig {
multi_scale: scales.len() > 1,
scales,
..Default::default()
};
group.bench_with_input(
BenchmarkId::new("multiscale", config_name),
&(test_image.clone(), config),
|b, (image, cfg)| {
b.iter(|| {
segmentor
.segment(image, cfg)
.expect("Multi-scale segmentation failed")
});
},
);
}
group.finish();
}
/// Benchmark real-time segmentation
fn bench_realtime_segmentation(c: &mut Criterion) {
let mut group = c.benchmark_group("realtime_segmentation");
// Test different input resolutions for real-time performance
let resolutions = vec![
("240p", 240, 320),
("360p", 360, 480),
("480p", 480, 640),
("720p", 720, 1280),
];
let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
.expect("Failed to create segmentor");
for (res_name, height, width) in resolutions {
let test_image = SegmentationBenchHelper::create_image(1, height, width);
let config = segmentation::SegmentationConfig {
vision_config: VisionConfig {
input_size: (height, width),
..Default::default()
},
output_stride: 16, // Larger stride for speed
..Default::default()
};
group.throughput(Throughput::Elements((height * width) as u64));
group.bench_with_input(
BenchmarkId::new("realtime", res_name),
&(test_image, config),
|b, (image, cfg)| {
b.iter(|| {
segmentor
.segment(image, cfg)
.expect("Real-time segmentation failed")
});
},
);
}
group.finish();
}
/// Benchmark medical image segmentation
fn bench_medical_segmentation(c: &mut Criterion) {
let mut group = c.benchmark_group("medical_segmentation");
// Different medical volume sizes
let volume_sizes = vec![
("ct_small", 32, 256, 256),
("ct_medium", 64, 512, 512),
("mri_large", 128, 256, 256),
];
for (volume_name, depth, height, width) in volume_sizes {
let medical_volume = SegmentationBenchHelper::create_medical_volume(depth, height, width);
// Medical segmentation typically processes slice by slice
let voxels = (depth * height * width) as u64;
group.throughput(Throughput::Elements(voxels));
group.bench_with_input(
BenchmarkId::new("medical_3d", volume_name),
&medical_volume,
|b, volume| {
b.iter(|| {
// Simulate medical volume segmentation
// In practice, this would involve:
// 1. Slice extraction
// 2. 2D segmentation on each slice
// 3. 3D reconstruction
// 4. Post-processing (morphological operations)
for slice_idx in 0..volume.shape()[1] {
let slice = volume
.narrow(1, slice_idx, 1)
.expect("Failed to extract slice");
let slice_2d = slice.squeeze_dim(1).expect("Failed to squeeze slice");
let slice_rgb = slice_2d
.unsqueeze_dim(1)
.expect("Failed to add channel")
.repeat(&[1, 3, 1, 1])
.expect("Failed to repeat channels");
// This is a placeholder for actual medical segmentation
let _ = slice_rgb;
}
});
},
);
}
group.finish();
}
/// Benchmark batch segmentation processing
fn bench_batch_segmentation(c: &mut Criterion) {
let mut group = c.benchmark_group("batch_segmentation");
let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
.expect("Failed to create segmentor");
let batch_sizes = vec![1, 2, 4, 8];
for batch_size in batch_sizes {
let batch_images = SegmentationBenchHelper::create_image(batch_size, 512, 512);
let config = segmentation::SegmentationConfig::default();
group.throughput(Throughput::Elements(batch_size as u64));
group.bench_with_input(
BenchmarkId::new("batch_process", batch_size),
&(batch_images, config),
|b, (images, cfg)| {
b.iter(|| {
segmentor
.segment_batch(images, cfg)
.expect("Batch segmentation failed")
});
},
);
}
group.finish();
}
/// Benchmark post-processing operations
fn bench_postprocessing(c: &mut Criterion) {
let mut group = c.benchmark_group("segmentation_postprocessing");
let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
.expect("Failed to create segmentor");
let test_image = SegmentationBenchHelper::create_image(1, 512, 512);
// Different post-processing configurations
let postprocess_configs = vec![
("no_postprocess", false, false, 0),
("morphology_only", true, false, 0),
("crf_only", false, true, 10),
("full_postprocess", true, true, 10),
];
for (config_name, use_morphology, use_crf, crf_iterations) in postprocess_configs {
let config = segmentation::SegmentationConfig {
post_process: segmentation::PostProcessConfig {
use_morphology,
use_crf,
crf_iterations,
..Default::default()
},
..Default::default()
};
group.bench_with_input(
BenchmarkId::new("postprocess", config_name),
&(test_image.clone(), config),
|b, (image, cfg)| {
b.iter(|| {
segmentor
.segment(image, cfg)
.expect("Post-processing benchmark failed")
});
},
);
}
group.finish();
}
/// Benchmark segmentation accuracy vs speed trade-offs
fn bench_accuracy_speed_tradeoff(c: &mut Criterion) {
let mut group = c.benchmark_group("segmentation_tradeoff");
let test_image = SegmentationBenchHelper::create_image(1, 512, 512);
// Different configurations representing accuracy vs speed trade-offs
let tradeoff_configs = vec![
("speed_optimized", 16, false, vec![1.0]), // Large stride, no multi-scale
("balanced", 8, true, vec![0.75, 1.0, 1.25]), // Medium stride, some multi-scale
(
"accuracy_optimized",
4,
true,
vec![0.5, 0.75, 1.0, 1.25, 1.5],
), // Small stride, full multi-scale
];
for (config_name, output_stride, multi_scale, scales) in tradeoff_configs {
let mut segmentor =
segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet101")
.expect("Failed to create segmentor");
let config = segmentation::SegmentationConfig {
output_stride,
multi_scale,
scales,
..Default::default()
};
group.bench_with_input(
BenchmarkId::new("tradeoff", config_name),
&(test_image.clone(), config),
|b, (image, cfg)| {
b.iter(|| {
segmentor
.segment(image, cfg)
.expect("Tradeoff benchmark failed")
});
},
);
}
group.finish();
}
/// Benchmark memory usage during segmentation
fn bench_memory_usage(c: &mut Criterion) {
let mut group = c.benchmark_group("segmentation_memory");
// Test memory usage with different image sizes and batch sizes
let test_configs = vec![
("small_single", 1, 256, 256),
("medium_single", 1, 512, 512),
("large_single", 1, 1024, 1024),
("small_batch", 4, 256, 256),
("medium_batch", 4, 512, 512),
];
for (config_name, batch_size, height, width) in test_configs {
let mut segmentor = segmentation::SegmentationFactory::create_deeplabv3plus(21, "resnet50")
.expect("Failed to create segmentor");
let test_images = SegmentationBenchHelper::create_image(batch_size, height, width);
let config = segmentation::SegmentationConfig {
vision_config: VisionConfig {
input_size: (height, width),
..Default::default()
},
..Default::default()
};
let total_pixels = (batch_size * height * width) as u64;
group.throughput(Throughput::Elements(total_pixels));
group.bench_with_input(
BenchmarkId::new("memory", config_name),
&(test_images, config),
|b, (images, cfg)| {
b.iter(|| {
if cfg.vision_config.input_size.0 > 512 {
// For large images, test single image processing
let single_image =
images.narrow(0, 0, 1).expect("Failed to get single image");
segmentor
.segment(&single_image, cfg)
.expect("Memory benchmark failed")
} else {
// For smaller images, test batch processing
segmentor
.segment_batch(images, cfg)
.expect("Memory benchmark failed")
}
});
},
);
}
group.finish();
}
criterion_group!(
benches,
bench_semantic_segmentation,
bench_instance_segmentation,
bench_multiscale_segmentation,
bench_realtime_segmentation,
bench_medical_segmentation,
bench_batch_segmentation,
bench_postprocessing,
bench_accuracy_speed_tradeoff,
bench_memory_usage
);
criterion_main!(benches);