# 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`:

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

## Quick Start

### Random Search

```rust
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

```rust
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:

```rust
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:

```rust
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:

```rust
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:

```rust
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:

```bash
cargo test -p rtx-nas
```

Run with output:

```bash
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](https://arxiv.org/abs/1806.09055)
- Neural Architecture Search: A Survey ([arXiv:1808.05377](https://arxiv.org/abs/1808.05377))