265 lines
8.8 KiB
Rust
265 lines
8.8 KiB
Rust
//! Hardware-Aware Neural Architecture Search Example
|
|
//!
|
|
//! This example demonstrates how to perform hardware-aware NAS using rtx-nas,
|
|
//! including:
|
|
//!
|
|
//! - Device profiling and latency prediction
|
|
//! - Multi-objective optimization (accuracy vs latency/memory)
|
|
//! - Pareto frontier construction
|
|
//! - FairNAS constraints for balanced architecture evaluation
|
|
//!
|
|
//! # Running
|
|
//!
|
|
//! ```bash
|
|
//! cargo run --example hardware_aware_search -p rtx-nas
|
|
//! ```
|
|
|
|
use rtx_nas::{
|
|
Result,
|
|
algorithms::{
|
|
PCDARTS, PCDARTSConfig, RandomSearch, RandomSearchConfig,
|
|
fairness::{FairnessConfig, FairnessTracker},
|
|
},
|
|
hardware::{
|
|
cost_model::compute_cost,
|
|
device::{CommonDevices, DeviceProfile},
|
|
latency::{LatencyPredictor, LookupTablePredictor},
|
|
},
|
|
search::{MultiObjective, ObjectiveScorer, ParetoEntry, ParetoFrontier},
|
|
search_space::{CellConfig, DARTSSearchSpace, SearchSpace},
|
|
};
|
|
use rtx_tensor::Device;
|
|
use std::time::Instant;
|
|
|
|
fn main() -> Result<()> {
|
|
println!("=== Hardware-Aware Neural Architecture Search ===\n");
|
|
|
|
// 1. Select target devices
|
|
println!("Step 1: Setting up target devices...");
|
|
let devices = vec![
|
|
CommonDevices::rtx_3090(),
|
|
CommonDevices::t4(),
|
|
CommonDevices::mobile_arm(),
|
|
];
|
|
|
|
for device in &devices {
|
|
println!(
|
|
" - {}: {:.1} TFLOPS, {:.0} GB memory",
|
|
device.name, device.peak_tflops_fp32, device.memory_gb
|
|
);
|
|
}
|
|
println!();
|
|
|
|
// 2. Create search space
|
|
println!("Step 2: Creating DARTS search space...");
|
|
let search_space = DARTSSearchSpace::default()?;
|
|
println!(
|
|
" Search space: {} total choices\n",
|
|
search_space.num_choices()
|
|
);
|
|
|
|
// 3. Initialize latency predictor
|
|
println!("Step 3: Setting up latency prediction...");
|
|
let predictor = LookupTablePredictor::new();
|
|
println!(
|
|
" Predictor confidence: {:.0}%\n",
|
|
predictor.confidence() * 100.0
|
|
);
|
|
|
|
// 4. Configure multi-objective optimization
|
|
println!("Step 4: Configuring multi-objective optimization...");
|
|
let objectives = MultiObjective::mobile_optimized();
|
|
println!(" Accuracy weight: {:.2}", objectives.accuracy_weight);
|
|
println!(" Latency weight: {:.2}", objectives.latency_weight);
|
|
println!(" Memory weight: {:.2}", objectives.memory_weight);
|
|
println!(" Params weight: {:.2}\n", objectives.params_weight);
|
|
|
|
let scorer = ObjectiveScorer::new(objectives)?;
|
|
|
|
// 5. Initialize fairness tracking
|
|
println!("Step 5: Setting up fairness tracking...");
|
|
let fairness_config = FairnessConfig::default()
|
|
.with_tracking_window(500)
|
|
.with_threshold(0.85);
|
|
let mut fairness_tracker = FairnessTracker::new(fairness_config);
|
|
println!(" Fairness threshold: 85%\n");
|
|
|
|
// 6. Sample architectures using random search
|
|
println!("Step 6: Sampling architectures with random search...");
|
|
let start = Instant::now();
|
|
|
|
let config = RandomSearchConfig::new(20);
|
|
let mut search = RandomSearch::new(config)?;
|
|
search.sample(&search_space)?;
|
|
|
|
println!(
|
|
" Sampled {} architectures in {:?}\n",
|
|
search.num_samples(),
|
|
start.elapsed()
|
|
);
|
|
|
|
// 7. Evaluate architectures and build Pareto frontier
|
|
println!("Step 7: Evaluating architectures...\n");
|
|
let mut frontier = ParetoFrontier::with_max_size(10);
|
|
|
|
let primary_device = &devices[0]; // Use RTX 3090 as primary evaluation device
|
|
|
|
for (i, arch) in search.samples().iter().enumerate() {
|
|
// Compute cost
|
|
let mut cost = compute_cost(arch)?;
|
|
|
|
// Predict latency
|
|
let latency = predictor.predict(arch, primary_device)?;
|
|
cost.estimated_latency_ms = Some(latency);
|
|
|
|
// Simulate accuracy (in practice, this comes from actual training)
|
|
let accuracy = simulate_accuracy(arch.channels, arch.num_layers);
|
|
|
|
// Track fairness (simulate operation usage from the architecture)
|
|
for cell in &arch.cells {
|
|
for edge in cell.edges() {
|
|
if let Some(op_type) = cell.get_operation(&edge) {
|
|
fairness_tracker.track_optimization(edge, op_type);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Compute weighted score
|
|
let score = scorer.score(accuracy, &cost);
|
|
|
|
// Add to Pareto frontier
|
|
let entry = ParetoEntry::new(arch.clone(), cost.clone(), accuracy);
|
|
let added = frontier.add(entry);
|
|
|
|
println!(
|
|
" Arch {:2}: acc={:.3}, latency={:.2}ms, params={:>6}, FLOPs={:>8}, score={:.3} {}",
|
|
i + 1,
|
|
accuracy,
|
|
latency,
|
|
format_number(cost.params),
|
|
format_number(cost.flops),
|
|
score,
|
|
if added { "[Pareto]" } else { "" }
|
|
);
|
|
}
|
|
|
|
// 8. Display Pareto frontier
|
|
println!("\n=== Pareto Frontier ===");
|
|
println!("Non-dominated architectures that offer best tradeoffs:\n");
|
|
|
|
for (i, entry) in frontier.entries().iter().enumerate() {
|
|
println!(
|
|
" {}. {} - Accuracy: {:.3}, Latency: {:.2}ms, Params: {}",
|
|
i + 1,
|
|
entry.arch_id(),
|
|
entry.accuracy,
|
|
entry.cost.estimated_latency_ms.unwrap_or(0.0),
|
|
format_number(entry.cost.params)
|
|
);
|
|
}
|
|
|
|
// 9. Display fairness report
|
|
println!("\n=== Fairness Analysis ===");
|
|
let report = fairness_tracker.generate_report();
|
|
println!(
|
|
" Overall fairness score: {:.1}%",
|
|
report.overall_score * 100.0
|
|
);
|
|
println!(" Total operations tracked: {}", report.total_iterations);
|
|
|
|
if report.underrepresented.is_empty() {
|
|
println!(" All operations fairly represented!");
|
|
} else {
|
|
println!(" Underrepresented operations:");
|
|
for op in &report.underrepresented {
|
|
println!(" - {:?}", op);
|
|
}
|
|
}
|
|
|
|
// 10. Demonstrate PC-DARTS with hardware awareness
|
|
println!("\n=== PC-DARTS with Hardware Constraints ===\n");
|
|
demonstrate_pcdarts(primary_device, &predictor)?;
|
|
|
|
println!("\n=== Hardware-Aware Search Complete ===");
|
|
Ok(())
|
|
}
|
|
|
|
/// Simulate accuracy based on architecture complexity
|
|
fn simulate_accuracy(channels: usize, num_layers: usize) -> f32 {
|
|
// Simple heuristic: more channels and layers = higher accuracy (with noise)
|
|
let base = 0.7 + (channels as f32 / 100.0).min(0.15) + (num_layers as f32 / 50.0).min(0.1);
|
|
let noise = (rand::random::<f32>() - 0.5) * 0.05;
|
|
(base + noise).clamp(0.5, 0.98)
|
|
}
|
|
|
|
/// Format large numbers for display
|
|
fn format_number(n: u64) -> String {
|
|
if n >= 1_000_000_000 {
|
|
format!("{:.1}B", n as f64 / 1_000_000_000.0)
|
|
} else if n >= 1_000_000 {
|
|
format!("{:.1}M", n as f64 / 1_000_000.0)
|
|
} else if n >= 1_000 {
|
|
format!("{:.1}K", n as f64 / 1_000.0)
|
|
} else {
|
|
n.to_string()
|
|
}
|
|
}
|
|
|
|
/// Demonstrate PC-DARTS with hardware constraints
|
|
fn demonstrate_pcdarts(device: &DeviceProfile, predictor: &LookupTablePredictor) -> Result<()> {
|
|
let compute_device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
println!("Creating PC-DARTS with partial channel connections...");
|
|
println!(" Target device: {}", device.name);
|
|
println!(" Channel fraction: 12.5% (memory efficient)");
|
|
println!(" Edge normalization: enabled\n");
|
|
|
|
// Create PC-DARTS configuration
|
|
let mut config = PCDARTSConfig::default();
|
|
config.edge_normalization = true;
|
|
|
|
let cell_configs = vec![CellConfig::default_darts()];
|
|
let mut pcdarts = PCDARTS::new(config, cell_configs, &compute_device)?;
|
|
|
|
// Simulate a few search steps
|
|
println!("Running architecture search steps...");
|
|
for epoch in 0..3 {
|
|
// In practice, train_loss and valid_loss come from actual training
|
|
let train_loss = 0.5 - (epoch as f32 * 0.1);
|
|
let valid_loss = 0.6 - (epoch as f32 * 0.08);
|
|
|
|
pcdarts.step(train_loss, valid_loss, None, None)?;
|
|
|
|
// Derive architecture and estimate cost
|
|
let arch = pcdarts.derive_architecture()?;
|
|
let cost = compute_cost(&arch)?;
|
|
let latency = predictor.predict(&arch, device)?;
|
|
|
|
println!(
|
|
" Epoch {}: train_loss={:.3}, valid_loss={:.3}, est_latency={:.2}ms",
|
|
epoch + 1,
|
|
train_loss,
|
|
valid_loss,
|
|
latency
|
|
);
|
|
}
|
|
|
|
// Derive final architecture
|
|
let final_arch = pcdarts.derive_architecture()?;
|
|
let final_cost = compute_cost(&final_arch)?;
|
|
let final_latency = predictor.predict(&final_arch, device)?;
|
|
|
|
println!("\nDerived architecture:");
|
|
println!(" ID: {}", final_arch.id);
|
|
println!(" Cells: {}", final_arch.num_cells());
|
|
println!(" Channels: {}", final_arch.channels);
|
|
println!(" Estimated FLOPs: {}", format_number(final_cost.flops));
|
|
println!(" Estimated params: {}", format_number(final_cost.params));
|
|
println!(
|
|
" Estimated latency on {}: {:.2}ms",
|
|
device.name, final_latency
|
|
);
|
|
|
|
Ok(())
|
|
}
|