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

737 lines
21 KiB
Rust

//! Simplified RegNet tests - TDD Red Phase
//!
//! This file tests RegNet implementation in isolation using minimal mock tensor framework.
use std::fmt;
/// Mock tensor error for testing
#[derive(Debug)]
pub enum MockTensorError {
ShapeMismatch(String),
DeviceMismatch(String),
InvalidOperation(String),
}
impl fmt::Display for MockTensorError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
MockTensorError::ShapeMismatch(msg) => write!(f, "Shape mismatch: {}", msg),
MockTensorError::DeviceMismatch(msg) => write!(f, "Device mismatch: {}", msg),
MockTensorError::InvalidOperation(msg) => write!(f, "Invalid operation: {}", msg),
}
}
}
impl std::error::Error for MockTensorError {}
pub type Result<T> = std::result::Result<T, MockTensorError>;
/// Mock device for testing
#[derive(Debug, Clone, PartialEq)]
pub enum Device {
Cpu,
}
impl Device {
pub fn cpu() -> Self {
Self::Cpu
}
}
/// Mock tensor for testing
#[derive(Debug, Clone)]
pub struct Tensor {
data: Vec<f32>,
shape: Vec<usize>,
}
impl Tensor {
pub fn randn(shape: Vec<usize>, _device: &Device) -> Result<Self> {
let numel: usize = shape.iter().product();
let data = vec![0.5; numel]; // Simplified - use constant instead of random
Ok(Self { data, shape })
}
pub fn shape(&self) -> &[usize] {
&self.shape
}
}
// ========== VISION ERROR TYPE ==========
#[derive(Debug)]
pub enum VisionError {
InvalidDimensions { expected: String, got: String },
ConfigurationError(String),
TensorError(String),
InvalidInput(String),
}
impl fmt::Display for VisionError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
VisionError::InvalidDimensions { expected, got } =>
write!(f, "Invalid dimensions: expected {}, got {}", expected, got),
VisionError::ConfigurationError(msg) => write!(f, "Configuration error: {}", msg),
VisionError::TensorError(msg) => write!(f, "Tensor operation error: {}", msg),
VisionError::InvalidInput(msg) => write!(f, "Invalid input: {}", msg),
}
}
}
impl std::error::Error for VisionError {}
// Convert MockTensorError to VisionError
impl From<MockTensorError> for VisionError {
fn from(err: MockTensorError) -> Self {
match err {
MockTensorError::ShapeMismatch(msg) => VisionError::InvalidDimensions {
expected: "compatible shape".to_string(),
got: msg
},
MockTensorError::DeviceMismatch(msg) => VisionError::TensorError(msg),
MockTensorError::InvalidOperation(msg) => VisionError::TensorError(msg),
}
}
}
pub type VisionResult<T> = std::result::Result<T, VisionError>;
// ========== REGNET CONFIGURATION ==========
#[derive(Clone, Debug)]
pub struct RegNetConfig {
w_a: f32,
w_0: f32,
w_m: f32,
group_width: usize,
depth_multiplier: f32,
num_classes: usize,
use_se: bool,
initial_width: f32,
num_stages: usize,
width_multiplier: f32,
}
#[derive(Clone, Debug)]
pub struct RegNetBlockConfig {
in_channels: usize,
out_channels: usize,
stride: usize,
groups: usize,
use_se: bool,
}
#[derive(Clone, Debug)]
pub struct SEModuleConfig {
channels: usize,
reduction: usize,
}
#[derive(Clone, Debug)]
pub struct RegNetStageConfig {
in_channels: usize,
out_channels: usize,
num_blocks: usize,
groups: usize,
stride: usize,
use_se: bool,
}
// ========== REGNET COMPONENTS ==========
pub struct RegNet {
stem: RegNetStem,
stages: Vec<RegNetStage>,
head: RegNetHead,
config: RegNetConfig,
}
pub struct RegNetBlock {
config: RegNetBlockConfig,
conv1_weight: Tensor,
conv2_weight: Tensor,
conv3_weight: Tensor,
bn1_weight: Tensor,
bn2_weight: Tensor,
bn3_weight: Tensor,
se_module: Option<SEModule>,
}
pub struct SEModule {
config: SEModuleConfig,
fc1_weight: Tensor,
fc2_weight: Tensor,
}
pub struct RegNetStem {
conv_weight: Tensor,
bn_weight: Tensor,
}
pub struct RegNetStage {
config: RegNetStageConfig,
blocks: Vec<RegNetBlock>,
}
pub struct RegNetHead {
pool_kernel_size: usize,
fc_weight: Tensor,
fc_bias: Tensor,
}
// ========== PLACEHOLDER IMPLEMENTATIONS (WILL FAIL) ==========
impl RegNetConfig {
pub fn new(
w_a: f32,
w_0: f32,
w_m: f32,
group_width: usize,
depth_multiplier: f32,
num_classes: usize,
) -> VisionResult<Self> {
Ok(Self {
w_a,
w_0,
w_m,
group_width,
depth_multiplier,
num_classes,
use_se: false,
initial_width: w_0,
num_stages: 4,
width_multiplier: w_m,
})
}
pub fn regnetx_200mf() -> Self {
Self {
w_a: 36.44,
w_0: 24.0,
w_m: 2.24,
group_width: 8,
depth_multiplier: 1.0,
num_classes: 1000,
use_se: false,
initial_width: 24.0,
num_stages: 4,
width_multiplier: 2.24,
}
}
pub fn regnety_200mf() -> Self {
Self {
w_a: 36.44,
w_0: 24.0,
w_m: 2.24,
group_width: 8,
depth_multiplier: 1.0,
num_classes: 1000,
use_se: true,
initial_width: 24.0,
num_stages: 4,
width_multiplier: 2.24,
}
}
pub fn regnetx_400mf() -> Self {
Self {
w_a: 24.48,
w_0: 24.0,
w_m: 2.54,
group_width: 16,
depth_multiplier: 1.0,
num_classes: 1000,
use_se: false,
initial_width: 24.0,
num_stages: 4,
width_multiplier: 2.54,
}
}
pub fn regnety_400mf() -> Self {
Self {
w_a: 27.89,
w_0: 48.0,
w_m: 2.09,
group_width: 8,
depth_multiplier: 1.0,
num_classes: 1000,
use_se: true,
initial_width: 48.0,
num_stages: 4,
width_multiplier: 2.09,
}
}
pub fn regnetx_600mf() -> Self {
Self {
w_a: 36.97,
w_0: 48.0,
w_m: 2.24,
group_width: 24,
depth_multiplier: 1.0,
num_classes: 1000,
use_se: false,
initial_width: 48.0,
num_stages: 4,
width_multiplier: 2.24,
}
}
pub fn regnetx_800mf() -> Self {
Self {
w_a: 35.73,
w_0: 56.0,
w_m: 2.28,
group_width: 16,
depth_multiplier: 1.0,
num_classes: 1000,
use_se: false,
initial_width: 56.0,
num_stages: 4,
width_multiplier: 2.28,
}
}
pub fn w_a(&self) -> f32 { self.w_a }
pub fn w_0(&self) -> f32 { self.w_0 }
pub fn w_m(&self) -> f32 { self.w_m }
pub fn group_width(&self) -> usize { self.group_width }
pub fn depth_multiplier(&self) -> f32 { self.depth_multiplier }
pub fn num_classes(&self) -> usize { self.num_classes }
pub fn use_se(&self) -> bool { self.use_se }
pub fn calculate_stage_widths(&self) -> Vec<usize> {
// Calculate stage widths based on RegNet scaling rules
let mut widths = Vec::new();
let mut current_width = self.initial_width as f32;
for _ in 0..self.num_stages {
widths.push((current_width as usize).max(8)); // Minimum width of 8
current_width *= self.width_multiplier;
}
// Round to nearest multiple of group_width for group convolutions
widths.iter().map(|&w| {
let remainder = w % self.group_width;
if remainder == 0 {
w
} else {
w + (self.group_width - remainder)
}
}).collect()
}
pub fn calculate_stage_depths(&self) -> Vec<usize> {
// Calculate stage depths based on depth multiplier
let base_depths = [2, 3, 4, 2]; // Common RegNet depth pattern
base_depths.iter()
.take(self.num_stages)
.map(|&d| {
let scaled_depth = (d as f32 * self.depth_multiplier).round() as usize;
scaled_depth.max(1) // Minimum depth of 1
})
.collect()
}
}
impl RegNet {
pub fn new(config: &RegNetConfig, device: &Device) -> VisionResult<Self> {
// Calculate stage configurations
let stage_widths = config.calculate_stage_widths();
let stage_depths = config.calculate_stage_depths();
// Create stem (initial conv layer)
let stem = RegNetStem {
conv_weight: Tensor::randn(vec![64, 3, 3, 3], device)?,
bn_weight: Tensor::randn(vec![64], device)?,
};
// Create stages
let mut stages = Vec::new();
let mut in_channels = 64;
for (stage_idx, (&out_channels, &depth)) in stage_widths.iter().zip(stage_depths.iter()).enumerate() {
let stride = if stage_idx == 0 { 1 } else { 2 };
let stage_config = RegNetStageConfig {
in_channels,
out_channels,
num_blocks: depth,
groups: config.group_width,
stride,
use_se: config.use_se,
};
let stage = RegNetStage::new(&stage_config, device)?;
stages.push(stage);
in_channels = out_channels;
}
// Create classification head
let head = RegNetHead {
pool_kernel_size: 7,
fc_weight: Tensor::randn(vec![config.num_classes, in_channels], device)?,
fc_bias: Tensor::randn(vec![config.num_classes], device)?,
};
Ok(Self {
stem,
stages,
head,
config: config.clone(),
})
}
pub fn forward(&self, input: &Tensor) -> VisionResult<Tensor> {
// Validate input shape (B, C, H, W)
if input.shape().len() != 4 {
return Err(VisionError::InvalidInput(
format!("Expected 4D input tensor (B,C,H,W), got {}D", input.shape().len())
));
}
if input.shape()[1] != 3 {
return Err(VisionError::InvalidInput(
format!("Expected 3 input channels, got {}", input.shape()[1])
));
}
// Forward through stem
let mut x = self.stem.forward(input)?;
// Forward through stages
for stage in &self.stages {
x = stage.forward(&x)?;
}
// Forward through head
x = self.head.forward(&x)?;
Ok(x)
}
}
impl RegNetBlockConfig {
pub fn new(in_channels: usize, out_channels: usize, stride: usize, groups: usize) -> Self {
Self {
in_channels,
out_channels,
stride,
groups,
use_se: false,
}
}
pub fn groups(&self) -> usize { self.groups }
pub fn out_channels(&self) -> usize { self.out_channels }
}
impl RegNetBlock {
pub fn new(config: &RegNetBlockConfig, device: &Device) -> VisionResult<Self> {
let se_module = if config.use_se {
let se_config = SEModuleConfig::new(config.out_channels, 4);
Some(SEModule::new(&se_config, device)?)
} else {
None
};
Ok(Self {
config: config.clone(),
conv1_weight: Tensor::randn(vec![config.out_channels, config.in_channels, 1, 1], device)?,
conv2_weight: Tensor::randn(vec![config.out_channels, config.out_channels / config.groups, 3, 3], device)?,
conv3_weight: Tensor::randn(vec![config.out_channels, config.out_channels, 1, 1], device)?,
bn1_weight: Tensor::randn(vec![config.out_channels], device)?,
bn2_weight: Tensor::randn(vec![config.out_channels], device)?,
bn3_weight: Tensor::randn(vec![config.out_channels], device)?,
se_module,
})
}
pub fn forward(&self, input: &Tensor) -> VisionResult<Tensor> {
// Simplified forward pass for testing
// In real implementation, this would do:
// 1. 1x1 conv + BN + ReLU
// 2. 3x3 grouped conv + BN + ReLU
// 3. 1x1 conv + BN
// 4. SE module (if enabled)
// 5. Residual connection + ReLU
let output_shape = vec![
input.shape()[0], // batch
self.config.out_channels, // channels
input.shape()[2] / self.config.stride, // height
input.shape()[3] / self.config.stride, // width
];
Ok(Tensor::randn(output_shape, &Device::cpu())?)
}
}
impl SEModuleConfig {
pub fn new(channels: usize, reduction: usize) -> Self {
Self { channels, reduction }
}
}
impl SEModule {
pub fn new(config: &SEModuleConfig, device: &Device) -> VisionResult<Self> {
let reduced_channels = config.channels / config.reduction;
Ok(Self {
config: config.clone(),
fc1_weight: Tensor::randn(vec![reduced_channels, config.channels], device)?,
fc2_weight: Tensor::randn(vec![config.channels, reduced_channels], device)?,
})
}
pub fn forward(&self, input: &Tensor) -> VisionResult<Tensor> {
// Simplified SE forward pass:
// 1. Global average pooling
// 2. FC1 + ReLU
// 3. FC2 + Sigmoid
// 4. Scale input
// For testing, just return input unchanged
Ok(input.clone())
}
}
impl RegNetStageConfig {
pub fn new(in_channels: usize, out_channels: usize, num_blocks: usize, groups: usize, stride: usize, use_se: bool) -> Self {
Self { in_channels, out_channels, num_blocks, groups, stride, use_se }
}
}
impl RegNetStage {
pub fn new(config: &RegNetStageConfig, device: &Device) -> VisionResult<Self> {
let mut blocks = Vec::new();
let mut in_channels = config.in_channels;
for i in 0..config.num_blocks {
let stride = if i == 0 { config.stride } else { 1 };
let block_config = RegNetBlockConfig {
in_channels,
out_channels: config.out_channels,
stride,
groups: config.groups,
use_se: config.use_se,
};
blocks.push(RegNetBlock::new(&block_config, device)?);
in_channels = config.out_channels;
}
Ok(Self {
config: config.clone(),
blocks,
})
}
pub fn forward(&self, input: &Tensor) -> VisionResult<Tensor> {
let mut x = input.clone();
for block in &self.blocks {
x = block.forward(&x)?;
}
Ok(x)
}
}
impl RegNetStem {
pub fn new(in_channels: usize, out_channels: usize, device: &Device) -> VisionResult<Self> {
Ok(Self {
conv_weight: Tensor::randn(vec![out_channels, in_channels, 3, 3], device)?,
bn_weight: Tensor::randn(vec![out_channels], device)?,
})
}
pub fn forward(&self, input: &Tensor) -> VisionResult<Tensor> {
// Simplified stem forward pass:
// 3x3 conv + BN + ReLU
let output_shape = vec![
input.shape()[0], // batch
self.conv_weight.shape()[0], // out_channels
input.shape()[2] / 2, // height (stride 2)
input.shape()[3] / 2, // width (stride 2)
];
Ok(Tensor::randn(output_shape, &Device::cpu())?)
}
}
impl RegNetHead {
pub fn new(in_channels: usize, num_classes: usize, device: &Device) -> VisionResult<Self> {
Ok(Self {
pool_kernel_size: 7,
fc_weight: Tensor::randn(vec![num_classes, in_channels], device)?,
fc_bias: Tensor::randn(vec![num_classes], device)?,
})
}
pub fn forward(&self, input: &Tensor) -> VisionResult<Tensor> {
// Simplified head forward pass:
// 1. Global average pooling
// 2. Fully connected layer
let batch_size = input.shape()[0];
let num_classes = self.fc_weight.shape()[0];
Ok(Tensor::randn(vec![batch_size, num_classes], &Device::cpu())?)
}
}
// ========== TDD RED PHASE TESTS (ALL SHOULD PANIC) ==========
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_regnet_design_space_parameters() {
let config = RegNetConfig::new(
32.0, // w_a
2.25, // w_0
1.0, // w_m
8, // group_width
2.0, // depth_multiplier
1000, // num_classes
).expect("Failed to create RegNet config");
assert_eq!(config.w_a(), 32.0);
assert_eq!(config.w_0(), 2.25);
assert_eq!(config.w_m(), 1.0);
assert_eq!(config.group_width(), 8);
assert_eq!(config.depth_multiplier(), 2.0);
assert_eq!(config.num_classes(), 1000);
println!("✓ Basic RegNet config creation works");
}
#[test]
fn test_quantized_linear_parameterization() {
let config = RegNetConfig::regnetx_200mf();
let stage_widths = config.calculate_stage_widths();
assert!(!stage_widths.is_empty());
println!("✓ Stage widths calculated: {:?}", stage_widths);
}
#[test]
fn test_stage_depth_calculation() {
let config = RegNetConfig::regnetx_200mf();
let stage_depths = config.calculate_stage_depths();
assert!(!stage_depths.is_empty());
println!("✓ Stage depths calculated: {:?}", stage_depths);
}
#[test]
fn test_regnet_block_creation() {
let device = Device::cpu();
let config = RegNetBlockConfig::new(64, 128, 1, 8);
let block = RegNetBlock::new(&config, &device).unwrap();
println!("✓ RegNet block created successfully");
assert_eq!(block.config.in_channels, 64);
assert_eq!(block.config.out_channels, 128);
}
#[test]
fn test_squeeze_excitation_module() {
let device = Device::cpu();
let config = SEModuleConfig::new(128, 4);
let se_module = SEModule::new(&config, &device).unwrap();
println!("✓ SE module created successfully");
assert_eq!(se_module.config.channels, 128);
assert_eq!(se_module.config.reduction, 4);
}
#[test]
fn test_regnet_stem() {
let device = Device::cpu();
let stem = RegNetStem::new(3, 32, &device).unwrap();
println!("✓ RegNet stem created successfully");
assert_eq!(stem.conv_weight.shape(), &[32, 3, 3, 3]);
}
#[test]
fn test_regnet_stage() {
let device = Device::cpu();
let config = RegNetStageConfig::new(64, 128, 2, 8, 2, true);
let stage = RegNetStage::new(&config, &device).unwrap();
println!("✓ RegNet stage created successfully with {} blocks", stage.blocks.len());
assert_eq!(stage.blocks.len(), 2);
}
#[test]
fn test_regnet_head() {
let device = Device::cpu();
let head = RegNetHead::new(512, 1000, &device).unwrap();
println!("✓ RegNet head created successfully");
assert_eq!(head.fc_weight.shape(), &[1000, 512]);
}
#[test]
fn test_regnet_complete_model() {
let device = Device::cpu();
let config = RegNetConfig::regnetx_200mf();
let model = RegNet::new(&config, &device).unwrap();
println!("✓ RegNet model created successfully with {} stages", model.stages.len());
}
#[test]
fn test_regnet_multiple_scales_config() {
let scales = vec![
RegNetConfig::regnetx_200mf(),
RegNetConfig::regnetx_400mf(),
RegNetConfig::regnetx_600mf(),
RegNetConfig::regnetx_800mf(),
RegNetConfig::regnety_200mf(),
RegNetConfig::regnety_400mf(),
];
for config in scales {
assert!(config.w_a() > 0.0);
assert!(config.w_0() > 0.0);
assert!(config.w_m() > 0.0);
assert!(config.group_width() > 0);
assert_eq!(config.num_classes(), 1000);
}
println!("✓ All RegNet scale configurations are valid");
}
#[test]
fn test_regnet_x_vs_y_differences() {
let regnetx = RegNetConfig::regnetx_200mf();
let regnety = RegNetConfig::regnety_200mf();
assert!(!regnetx.use_se());
assert!(regnety.use_se());
println!("✓ RegNetX/Y SE difference validated");
}
#[test]
fn test_grouped_convolution_config() {
let config = RegNetBlockConfig::new(64, 128, 2, 8);
assert_eq!(config.groups(), 8);
assert_eq!(config.out_channels() % config.groups(), 0);
println!("✓ Grouped convolution configuration valid");
}
}
fn main() {
println!("=== RegNet TDD Red Phase ===");
println!("Running comprehensive failing tests...");
// Run a few basic tests manually to show they work
let config = RegNetConfig::regnetx_200mf();
println!("RegNetX-200MF: w_a={}, w_0={}, group_width={}, use_se={}",
config.w_a(), config.w_0(), config.group_width(), config.use_se());
let config_y = RegNetConfig::regnety_200mf();
println!("RegNetY-200MF: w_a={}, w_0={}, group_width={}, use_se={}",
config_y.w_a(), config_y.w_0(), config_y.group_width(), config_y.use_se());
println!("\nRed phase complete! All implementation methods should panic when called.");
println!("Next: Implement minimal functionality to make tests pass (Green phase)");
}