128 lines
3.2 KiB
Rust
128 lines
3.2 KiB
Rust
//! Neural Architecture Search (NAS) for RustyTorch
|
|
//!
|
|
//! This crate provides implementations of Neural Architecture Search algorithms,
|
|
//! including DARTS (Differentiable Architecture Search) and random search baselines.
|
|
//!
|
|
//! # Overview
|
|
//!
|
|
//! Neural Architecture Search automates the design of neural network architectures.
|
|
//! This crate provides:
|
|
//!
|
|
//! - **Search Spaces**: Define the space of possible architectures
|
|
//! - Cell-based search space (DARTS-style)
|
|
//! - Operation primitives (convolutions, pooling, etc.)
|
|
//!
|
|
//! - **Search Algorithms**:
|
|
//! - DARTS: Gradient-based architecture search
|
|
//! - Random Search: Baseline algorithm
|
|
//!
|
|
//! # Examples
|
|
//!
|
|
//! ## Random Search
|
|
//!
|
|
//! ```rust
|
|
//! use rtx_nas::{
|
|
//! algorithms::{RandomSearch, RandomSearchConfig},
|
|
//! search_space::DARTSSearchSpace,
|
|
//! };
|
|
//!
|
|
//! # fn main() -> rtx_nas::error::Result<()> {
|
|
//! // 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());
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
//!
|
|
//! ## DARTS Algorithm
|
|
//!
|
|
//! ```rust
|
|
//! use rtx_nas::{
|
|
//! algorithms::{DARTS, DARTSConfig},
|
|
//! search_space::CellConfig,
|
|
//! };
|
|
//! use rtx_tensor::Device;
|
|
//!
|
|
//! # fn main() -> rtx_nas::error::Result<()> {
|
|
//! 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()?;
|
|
//! println!("Found architecture with {} cells", architecture.num_cells());
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
//!
|
|
//! ## Working with Cells
|
|
//!
|
|
//! ```rust
|
|
//! use rtx_nas::search_space::{Cell, CellConfig, Edge, OperationType};
|
|
//!
|
|
//! # fn main() -> rtx_nas::error::Result<()> {
|
|
//! // Create a cell
|
|
//! 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 cell structure
|
|
//! println!("Total nodes: {}", cell.total_nodes());
|
|
//! println!("Edges: {:?}", cell.edges());
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
|
|
pub mod algorithms;
|
|
pub mod error;
|
|
pub mod hardware;
|
|
pub mod search;
|
|
pub mod search_space;
|
|
|
|
// Re-export commonly used types
|
|
pub use algorithms::{
|
|
// PC-DARTS
|
|
ChannelMask,
|
|
// DARTS
|
|
DARTS,
|
|
DARTSCell,
|
|
DARTSConfig,
|
|
// FairNAS
|
|
FairnessAware,
|
|
FairnessConfig,
|
|
FairnessReport,
|
|
FairnessTracker,
|
|
MixedOp,
|
|
PCDARTS,
|
|
PCDARTSCell,
|
|
PCDARTSConfig,
|
|
PartialChannelMixedOp,
|
|
// Random Search
|
|
RandomSearch,
|
|
RandomSearchConfig,
|
|
};
|
|
pub use error::{NASError, Result};
|
|
pub use search_space::{
|
|
Architecture, Cell, CellConfig, DARTSSearchSpace, Edge, Operation, OperationConfig,
|
|
OperationType, SearchSpace,
|
|
};
|