214 lines
5.0 KiB
Markdown
214 lines
5.0 KiB
Markdown
# 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))
|