189 lines
7.5 KiB
Rust
189 lines
7.5 KiB
Rust
use rtx_multimodal::fusion::{FusionConfig, ModalityFusion};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
|
|
#[cfg(test)]
|
|
mod fusion_tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_modality_fusion_creation() {
|
|
// RED: Write failing test first
|
|
let device = Device::default();
|
|
let config = FusionConfig::new(768, 768, 768, 768, 12);
|
|
|
|
let fusion = ModalityFusion::new(config, &device);
|
|
assert!(fusion.is_ok(), "Should create ModalityFusion successfully");
|
|
|
|
let fusion = fusion.unwrap();
|
|
assert_eq!(fusion.config().vision_dim, 768);
|
|
assert_eq!(fusion.config().audio_dim, 768);
|
|
assert_eq!(fusion.config().text_dim, 768);
|
|
assert_eq!(fusion.config().output_dim, 768);
|
|
assert_eq!(fusion.config().num_heads, 12);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bimodal_fusion_vision_audio() {
|
|
// RED: Test bimodal fusion
|
|
let device = Device::default();
|
|
let config = FusionConfig::new(512, 512, 512, 512, 8);
|
|
let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");
|
|
|
|
// Create test tensors
|
|
let batch_size = 2;
|
|
let seq_len_vision = 197; // 196 patches + 1 cls token
|
|
let seq_len_audio = 500;
|
|
let hidden_dim = 512;
|
|
|
|
let vision = Tensor::randn(&[batch_size, seq_len_vision, hidden_dim], &device)
|
|
.expect("Failed to create vision tensor");
|
|
let audio = Tensor::randn(&[batch_size, seq_len_audio, hidden_dim], &device)
|
|
.expect("Failed to create audio tensor");
|
|
|
|
// Create dummy text tensor for trimodal fusion
|
|
let text = Tensor::randn(&[batch_size, 50, hidden_dim], &device)
|
|
.expect("Failed to create text tensor");
|
|
let result = fusion.forward_trimodal(&vision, &audio, &text);
|
|
assert!(result.is_ok(), "Bimodal fusion should succeed");
|
|
|
|
let output = result.unwrap();
|
|
assert_eq!(output.shape()[0], batch_size);
|
|
assert_eq!(output.shape()[1], hidden_dim);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bimodal_fusion_vision_text() {
|
|
// RED: Test vision-text fusion
|
|
let device = Device::default();
|
|
let config = FusionConfig::new(384, 384, 384, 384, 6);
|
|
let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");
|
|
|
|
let batch_size = 4;
|
|
let seq_len_vision = 100;
|
|
let seq_len_text = 64;
|
|
let hidden_dim = 384;
|
|
|
|
let vision = Tensor::randn(&[batch_size, seq_len_vision, hidden_dim], &device)
|
|
.expect("Failed to create vision tensor");
|
|
let text = Tensor::randn(&[batch_size, seq_len_text, hidden_dim], &device)
|
|
.expect("Failed to create text tensor");
|
|
|
|
// Create dummy audio tensor for trimodal fusion
|
|
let audio = Tensor::randn(&[batch_size, 80, hidden_dim], &device)
|
|
.expect("Failed to create audio tensor");
|
|
let result = fusion.forward_trimodal(&vision, &audio, &text);
|
|
assert!(result.is_ok(), "Vision-text fusion should succeed");
|
|
|
|
let output = result.unwrap();
|
|
assert_eq!(output.shape()[0], batch_size);
|
|
assert_eq!(output.shape()[1], hidden_dim);
|
|
}
|
|
|
|
#[test]
|
|
fn test_trimodal_fusion() {
|
|
// RED: Test trimodal fusion
|
|
let device = Device::default();
|
|
let config = FusionConfig::new(256, 256, 256, 256, 4);
|
|
let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");
|
|
|
|
let batch_size = 1;
|
|
let vision =
|
|
Tensor::randn(&[batch_size, 50, 256], &device).expect("Failed to create vision tensor");
|
|
let audio =
|
|
Tensor::randn(&[batch_size, 100, 256], &device).expect("Failed to create audio tensor");
|
|
let text =
|
|
Tensor::randn(&[batch_size, 30, 256], &device).expect("Failed to create text tensor");
|
|
|
|
let result = fusion.forward_trimodal(&vision, &audio, &text);
|
|
assert!(result.is_ok(), "Trimodal fusion should succeed");
|
|
|
|
let output = result.unwrap();
|
|
assert_eq!(output.shape()[0], batch_size);
|
|
assert_eq!(output.shape()[1], 256); // output_dim
|
|
}
|
|
|
|
#[test]
|
|
fn test_fusion_strategies() {
|
|
// RED: Test different fusion strategies
|
|
use rtx_multimodal::fusion::ModalityFusionStrategy;
|
|
|
|
let device = Device::default();
|
|
let strategies = vec![
|
|
ModalityFusionStrategy::Concatenation,
|
|
ModalityFusionStrategy::BilinearFusion,
|
|
ModalityFusionStrategy::TensorFusion,
|
|
ModalityFusionStrategy::AttentionFusion,
|
|
ModalityFusionStrategy::HierarchicalAttention,
|
|
ModalityFusionStrategy::GatedFusion,
|
|
ModalityFusionStrategy::ContrastiveFusion,
|
|
];
|
|
|
|
for strategy in strategies {
|
|
let mut config = FusionConfig::new(128, 128, 128, 128, 2);
|
|
config.fusion_strategy = strategy.clone();
|
|
|
|
let mut fusion = ModalityFusion::new(config, &device).expect(&format!(
|
|
"Failed to create fusion with strategy {:?}",
|
|
strategy
|
|
));
|
|
|
|
let vision = Tensor::randn(&[1, 10, 128], &device).unwrap();
|
|
let audio = Tensor::randn(&[1, 20, 128], &device).unwrap();
|
|
let text = Tensor::randn(&[1, 15, 128], &device).unwrap();
|
|
|
|
let result = fusion.forward_trimodal(&vision, &audio, &text);
|
|
assert!(result.is_ok(), "Strategy {:?} should work", strategy);
|
|
|
|
let output = result.unwrap();
|
|
assert_eq!(output.shape()[0], 1);
|
|
assert_eq!(output.shape()[1], 128);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_fusion_output_consistency() {
|
|
// RED: Test that fusion output is deterministic with same input
|
|
let device = Device::default();
|
|
let config = FusionConfig::new(64, 64, 64, 64, 2);
|
|
let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");
|
|
|
|
// Use zeros for deterministic testing
|
|
let vision = Tensor::zeros(&[1, 5, 64], &device).expect("Failed to create vision tensor");
|
|
let audio = Tensor::zeros(&[1, 10, 64], &device).expect("Failed to create audio tensor");
|
|
let text = Tensor::zeros(&[1, 8, 64], &device).expect("Failed to create text tensor");
|
|
|
|
let result1 = fusion
|
|
.forward_trimodal(&vision, &audio, &text)
|
|
.expect("First forward pass failed");
|
|
let result2 = fusion
|
|
.forward_trimodal(&vision, &audio, &text)
|
|
.expect("Second forward pass failed");
|
|
|
|
// With zero inputs and no randomness, outputs should be deterministic
|
|
// This tests that the fusion mechanism itself is stable
|
|
assert_eq!(result1.shape(), result2.shape());
|
|
}
|
|
|
|
#[test]
|
|
fn test_empty_batch() {
|
|
// RED: Test edge case with batch size 0
|
|
let device = Device::default();
|
|
let config = FusionConfig::new(64, 64, 64, 64, 2);
|
|
let mut fusion = ModalityFusion::new(config, &device).expect("Failed to create fusion");
|
|
|
|
// Empty batch
|
|
let vision =
|
|
Tensor::zeros(&[0, 5, 64], &device).expect("Failed to create empty vision tensor");
|
|
let audio =
|
|
Tensor::zeros(&[0, 10, 64], &device).expect("Failed to create empty audio tensor");
|
|
let text = Tensor::zeros(&[0, 8, 64], &device).expect("Failed to create empty text tensor");
|
|
|
|
let result = fusion.forward_trimodal(&vision, &audio, &text);
|
|
|
|
// Should handle empty batch gracefully
|
|
assert!(result.is_ok(), "Should handle empty batch");
|
|
let output = result.unwrap();
|
|
assert_eq!(output.shape()[0], 0);
|
|
}
|
|
}
|