rustytorch/crates/training/rtx-nas

rtx-nas

Neural Architecture Search (NAS) for RustyTorch - Automated neural network architecture discovery.

Overview

rtx-nas provides implementations of Neural Architecture Search algorithms for automatically discovering optimal neural network architectures. The crate includes:

• Search Spaces: Define the space of possible architectures

  • Cell-based search space (DARTS-style)
  • Operation primitives (convolutions, pooling, etc.)

• Search Algorithms:

  • DARTS: Differentiable Architecture Search (gradient-based)
  • Random Search: Baseline algorithm

Features

• ✅ Cell-based search space with 9 operation types
• ✅ DARTS algorithm with bi-level optimization
• ✅ Random search baseline
• ✅ Architecture encoding/decoding
• ✅ Comprehensive test coverage (91 tests)
• ✅ Full documentation with examples
• ✅ Zero unsafe code

Installation

Add this to your Cargo.toml:

[dependencies]
rtx-nas = { path = "path/to/rtx-nas" }

Quick Start

Random Search

use rtx_nas::{
    algorithms::{RandomSearch, RandomSearchConfig},
    search_space::DARTSSearchSpace,
};

// Create search space
let search_space = DARTSSearchSpace::default()?;

// Configure random search
let config = RandomSearchConfig::new(10);
let mut search = RandomSearch::new(config)?;

// Sample architectures
search.sample(&search_space)?;

println!("Sampled {} architectures", search.num_samples());

DARTS Algorithm

use rtx_nas::{
    algorithms::{DARTS, DARTSConfig},
    search_space::CellConfig,
};
use rtx_tensor::Device;

let device = Device::cuda(0).unwrap_or(Device::default());

// Configure DARTS
let config = DARTSConfig::default();
let cell_configs = vec![CellConfig::default_darts()];

// Create DARTS instance
let mut darts = DARTS::new(config, cell_configs, &device)?;

// Perform optimization steps
darts.step(0.5, 0.6, None, None)?;

// Derive final architecture
let architecture = darts.derive_architecture()?;

Search Space

The search space defines the possible architectures that can be discovered:

Operation Types

• Identity: Skip connection
• Zero: No connection
• Conv3x3: 3×3 convolution
• Conv5x5: 5×5 convolution
• SepConv3x3: Separable 3×3 convolution
• SepConv5x5: Separable 5×5 convolution
• DilConv3x3: Dilated 3×3 convolution
• MaxPool3x3: 3×3 max pooling
• AvgPool3x3: 3×3 average pooling

Cell Structure

Cells are the building blocks of architectures:

use rtx_nas::search_space::{Cell, CellConfig, Edge, OperationType};

// Create a cell with default DARTS configuration
let config = CellConfig::default_darts();
let mut cell = Cell::new(config)?;

// Set operations on edges
let edge = Edge::new(0, 2);
cell.set_operation(edge, OperationType::Conv3x3)?;

// Query structure
println!("Total nodes: {}", cell.total_nodes());
println!("Number of edges: {}", cell.edges().len());

Algorithms

DARTS (Differentiable Architecture Search)

DARTS enables gradient-based architecture search by relaxing the discrete architecture search space to be continuous:

let config = DARTSConfig {
    learning_rate_arch: 0.001,
    learning_rate_weights: 0.01,
    num_epochs: 50,
    warmup_epochs: 15,
    temperature: 1.0,
};

The algorithm alternates between:

1. Training network weights on training data
2. Optimizing architecture parameters on validation data

Random Search

A baseline algorithm that randomly samples architectures:

let config = RandomSearchConfig::new(100)
    .with_seed(42); // Optional: for reproducibility

let mut search = RandomSearch::new(config)?;
search.sample(&search_space)?;

// Get best architecture based on scores
let scores = vec![...]; // Evaluation scores
let best = search.get_best(&scores)?;

Architecture Encoding

Architectures can be encoded/decoded for storage and analysis:

let search_space = DARTSSearchSpace::default()?;
let arch = search_space.sample()?;

// Encode to continuous representation
let encoding = search_space.encode(&arch)?;

// Decode back to architecture
let decoded = search_space.decode(&encoding)?;

Examples

See the integration tests for complete examples:

• tests/integration_tests.rs: Comprehensive workflow examples
• Library documentation: Run cargo doc --open -p rtx-nas

Testing

Run all tests:

cargo test -p rtx-nas

Run with output:

cargo test -p rtx-nas -- --nocapture

Performance

The crate is designed for efficiency:

• Zero-cost abstractions using Rust's type system
• Minimal allocations in hot paths
• Efficient softmax computation for architecture weights

Contributing

Contributions are welcome! Please ensure:

• All tests pass: cargo test -p rtx-nas
• Code is formatted: cargo fmt
• Clippy is happy: cargo clippy -p rtx-nas
• Documentation is updated

License

MIT OR Apache-2.0

References

• DARTS: Differentiable Architecture Search
• Neural Architecture Search: A Survey (arXiv:1808.05377)