Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,672 @@
|
||||
//! DARTS (Differentiable Architecture Search) implementation
|
||||
|
||||
use crate::error::{NASError, Result};
|
||||
use crate::search_space::{Architecture, Cell, CellConfig, Edge, OperationType};
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Configuration for DARTS algorithm
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DARTSConfig {
|
||||
/// Learning rate for architecture parameters
|
||||
pub learning_rate_arch: f32,
|
||||
/// Learning rate for network weights
|
||||
pub learning_rate_weights: f32,
|
||||
/// Number of training epochs
|
||||
pub num_epochs: usize,
|
||||
/// Number of warmup epochs (train weights only)
|
||||
pub warmup_epochs: usize,
|
||||
/// Temperature for softmax
|
||||
pub temperature: f32,
|
||||
}
|
||||
|
||||
impl DARTSConfig {
|
||||
/// Create a new DARTS configuration
|
||||
pub fn new(
|
||||
learning_rate_arch: f32,
|
||||
learning_rate_weights: f32,
|
||||
num_epochs: usize,
|
||||
warmup_epochs: usize,
|
||||
) -> Self {
|
||||
Self {
|
||||
learning_rate_arch,
|
||||
learning_rate_weights,
|
||||
num_epochs,
|
||||
warmup_epochs,
|
||||
temperature: 1.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create default configuration
|
||||
pub fn default() -> Self {
|
||||
Self {
|
||||
learning_rate_arch: 0.001,
|
||||
learning_rate_weights: 0.01,
|
||||
num_epochs: 50,
|
||||
warmup_epochs: 15,
|
||||
temperature: 1.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate configuration
|
||||
pub fn validate(&self) -> Result<()> {
|
||||
if self.learning_rate_arch <= 0.0 {
|
||||
return Err(NASError::InvalidConfig(
|
||||
"learning_rate_arch must be positive".into(),
|
||||
));
|
||||
}
|
||||
if self.learning_rate_weights <= 0.0 {
|
||||
return Err(NASError::InvalidConfig(
|
||||
"learning_rate_weights must be positive".into(),
|
||||
));
|
||||
}
|
||||
if self.num_epochs == 0 {
|
||||
return Err(NASError::InvalidConfig(
|
||||
"num_epochs must be greater than 0".into(),
|
||||
));
|
||||
}
|
||||
if self.warmup_epochs > self.num_epochs {
|
||||
return Err(NASError::InvalidConfig(
|
||||
"warmup_epochs cannot exceed num_epochs".into(),
|
||||
));
|
||||
}
|
||||
if self.temperature <= 0.0 {
|
||||
return Err(NASError::InvalidConfig(
|
||||
"temperature must be positive".into(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Set temperature
|
||||
pub fn with_temperature(mut self, temperature: f32) -> Self {
|
||||
self.temperature = temperature;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Mixed operation that combines multiple operations with learned weights
|
||||
#[derive(Debug)]
|
||||
pub struct MixedOp {
|
||||
/// Edge this mixed op is for
|
||||
edge: Edge,
|
||||
/// Number of input/output channels
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
/// Stride
|
||||
stride: usize,
|
||||
/// Device
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl MixedOp {
|
||||
/// Create a new mixed operation
|
||||
pub fn new(
|
||||
edge: Edge,
|
||||
in_channels: usize,
|
||||
out_channels: usize,
|
||||
stride: usize,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
if in_channels == 0 || out_channels == 0 {
|
||||
return Err(NASError::InvalidConfig("Channels must be positive".into()));
|
||||
}
|
||||
Ok(Self {
|
||||
edge,
|
||||
in_channels,
|
||||
out_channels,
|
||||
stride,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward pass with architecture weights (alphas)
|
||||
pub fn forward(&self, input: &Tensor, weights: &[f32]) -> Result<Tensor> {
|
||||
if weights.len() != OperationType::count() {
|
||||
return Err(NASError::OperationError(format!(
|
||||
"Expected {} weights, got {}",
|
||||
OperationType::count(),
|
||||
weights.len()
|
||||
)));
|
||||
}
|
||||
|
||||
use crate::search_space::operations::{Operation, OperationConfig};
|
||||
|
||||
let ops = OperationType::all();
|
||||
let mut outputs = Vec::new();
|
||||
|
||||
// Execute each operation
|
||||
for (i, &op_type) in ops.iter().enumerate() {
|
||||
let config =
|
||||
OperationConfig::new(op_type, self.in_channels, self.out_channels, self.stride);
|
||||
let op = Operation::new(config, &self.device)?;
|
||||
let output = op.forward(input)?;
|
||||
|
||||
// Weight the output
|
||||
let weight_scalar = weights[i];
|
||||
let weighted = output.mul_scalar(weight_scalar)?;
|
||||
outputs.push(weighted);
|
||||
}
|
||||
|
||||
// Sum all weighted outputs
|
||||
let mut result = outputs[0].clone();
|
||||
for output in outputs.iter().skip(1) {
|
||||
result = result.add(output)?;
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Get the edge
|
||||
pub fn edge(&self) -> Edge {
|
||||
self.edge
|
||||
}
|
||||
}
|
||||
|
||||
/// DARTS cell with architecture parameters
|
||||
#[derive(Debug)]
|
||||
pub struct DARTSCell {
|
||||
config: CellConfig,
|
||||
/// Architecture parameters (alphas) for each edge
|
||||
/// Maps edge to vector of weights for each operation
|
||||
alphas: HashMap<Edge, Vec<f32>>,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl DARTSCell {
|
||||
/// Create a new DARTS cell
|
||||
pub fn new(config: CellConfig, device: &Device) -> Result<Self> {
|
||||
config.validate()?;
|
||||
|
||||
let cell = Cell::new(config.clone())?;
|
||||
let mut alphas = HashMap::new();
|
||||
|
||||
// Initialize alphas uniformly for each edge
|
||||
let num_ops = OperationType::count();
|
||||
let initial_value = 1.0 / num_ops as f32;
|
||||
|
||||
for edge in cell.edges() {
|
||||
alphas.insert(edge, vec![initial_value; num_ops]);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
alphas,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Get architecture parameters
|
||||
pub fn alphas(&self) -> &HashMap<Edge, Vec<f32>> {
|
||||
&self.alphas
|
||||
}
|
||||
|
||||
/// Get mutable architecture parameters
|
||||
pub fn alphas_mut(&mut self) -> &mut HashMap<Edge, Vec<f32>> {
|
||||
&mut self.alphas
|
||||
}
|
||||
|
||||
/// Apply softmax to get operation probabilities
|
||||
pub fn get_probabilities(&self, temperature: f32) -> HashMap<Edge, Vec<f32>> {
|
||||
let mut probs = HashMap::new();
|
||||
|
||||
for (edge, alpha_vec) in &self.alphas {
|
||||
let softmax_probs = softmax(alpha_vec, temperature);
|
||||
probs.insert(*edge, softmax_probs);
|
||||
}
|
||||
|
||||
probs
|
||||
}
|
||||
|
||||
/// Derive discrete architecture from current alphas
|
||||
pub fn derive_architecture(&self) -> Result<Cell> {
|
||||
let mut cell = Cell::new(self.config.clone())?;
|
||||
|
||||
for (edge, alpha_vec) in &self.alphas {
|
||||
// Select operation with highest alpha
|
||||
let op_idx = alpha_vec
|
||||
.iter()
|
||||
.enumerate()
|
||||
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
|
||||
.map(|(idx, _)| idx)
|
||||
.ok_or_else(|| NASError::OperationError("Empty alpha vector".into()))?;
|
||||
|
||||
let op = OperationType::all()[op_idx];
|
||||
cell.set_operation(*edge, op)?;
|
||||
}
|
||||
|
||||
Ok(cell)
|
||||
}
|
||||
|
||||
/// Get cell configuration
|
||||
pub fn config(&self) -> &CellConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Get device
|
||||
pub fn device(&self) -> &Device {
|
||||
&self.device
|
||||
}
|
||||
}
|
||||
|
||||
/// Softmax function
|
||||
fn softmax(values: &[f32], temperature: f32) -> Vec<f32> {
|
||||
let scaled: Vec<f32> = values.iter().map(|&x| x / temperature).collect();
|
||||
|
||||
let max = scaled.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
let exp_values: Vec<f32> = scaled.iter().map(|&x| (x - max).exp()).collect();
|
||||
let sum: f32 = exp_values.iter().sum();
|
||||
|
||||
exp_values.iter().map(|&x| x / sum).collect()
|
||||
}
|
||||
|
||||
/// DARTS algorithm
|
||||
#[derive(Debug)]
|
||||
pub struct DARTS {
|
||||
config: DARTSConfig,
|
||||
cells: Vec<DARTSCell>,
|
||||
device: Device,
|
||||
current_epoch: usize,
|
||||
}
|
||||
|
||||
impl DARTS {
|
||||
/// Create a new DARTS instance
|
||||
pub fn new(
|
||||
config: DARTSConfig,
|
||||
cell_configs: Vec<CellConfig>,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
config.validate()?;
|
||||
|
||||
if cell_configs.is_empty() {
|
||||
return Err(NASError::InvalidConfig(
|
||||
"At least one cell configuration required".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let mut cells = Vec::new();
|
||||
for cell_config in cell_configs {
|
||||
cells.push(DARTSCell::new(cell_config, device)?);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
cells,
|
||||
device: device.clone(),
|
||||
current_epoch: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create default DARTS with 2 cells (normal and reduction)
|
||||
pub fn default(device: &Device) -> Result<Self> {
|
||||
let config = DARTSConfig::default();
|
||||
let cell_configs = vec![CellConfig::default_darts(), CellConfig::default_darts()];
|
||||
Self::new(config, cell_configs, device)
|
||||
}
|
||||
|
||||
/// Perform one step of bi-level optimization
|
||||
///
|
||||
/// 1. Update weights on training data
|
||||
/// 2. Update architecture parameters on validation data
|
||||
pub fn step(
|
||||
&mut self,
|
||||
_train_loss: f32,
|
||||
val_loss: f32,
|
||||
_train_gradients: Option<&HashMap<Edge, Vec<f32>>>,
|
||||
val_gradients: Option<&HashMap<Edge, Vec<f32>>>,
|
||||
) -> Result<()> {
|
||||
// During warmup, only update weights (simulated here)
|
||||
if self.current_epoch < self.config.warmup_epochs {
|
||||
self.current_epoch += 1;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Update architecture parameters using validation gradients
|
||||
if let Some(val_grads) = val_gradients {
|
||||
for cell in &mut self.cells {
|
||||
for (edge, grad_vec) in val_grads {
|
||||
if let Some(alpha_vec) = cell.alphas_mut().get_mut(edge) {
|
||||
// Gradient descent on alphas
|
||||
for (alpha, &grad) in alpha_vec.iter_mut().zip(grad_vec.iter()) {
|
||||
*alpha -= self.config.learning_rate_arch * grad;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Use numerical gradient estimation based on loss
|
||||
self.update_alphas_numerical(val_loss)?;
|
||||
}
|
||||
|
||||
self.current_epoch += 1;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Update alphas using numerical gradient estimation
|
||||
fn update_alphas_numerical(&mut self, loss: f32) -> Result<()> {
|
||||
let epsilon = 1e-3;
|
||||
|
||||
for cell in &mut self.cells {
|
||||
for alpha_vec in cell.alphas_mut().values_mut() {
|
||||
for alpha in alpha_vec.iter_mut() {
|
||||
// Simple gradient descent with perturbation
|
||||
let perturbation = epsilon * (0.5 - rand::random::<f32>());
|
||||
*alpha -= self.config.learning_rate_arch * loss * perturbation;
|
||||
|
||||
// Clamp to prevent extreme values
|
||||
*alpha = alpha.clamp(-10.0, 10.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Derive final architecture from learned alphas
|
||||
pub fn derive_architecture(&self) -> Result<Architecture> {
|
||||
let mut cells = Vec::new();
|
||||
|
||||
for darts_cell in &self.cells {
|
||||
cells.push(darts_cell.derive_architecture()?);
|
||||
}
|
||||
|
||||
let id = format!("darts_arch_{}", uuid::Uuid::new_v4().simple());
|
||||
Ok(Architecture::new(id, cells, 16, 8))
|
||||
}
|
||||
|
||||
/// Get current alphas for all cells
|
||||
pub fn get_alphas(&self) -> Vec<&HashMap<Edge, Vec<f32>>> {
|
||||
self.cells.iter().map(DARTSCell::alphas).collect()
|
||||
}
|
||||
|
||||
/// Get current architecture probabilities
|
||||
pub fn get_probabilities(&self) -> Vec<HashMap<Edge, Vec<f32>>> {
|
||||
self.cells
|
||||
.iter()
|
||||
.map(|cell| cell.get_probabilities(self.config.temperature))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get current epoch
|
||||
pub fn current_epoch(&self) -> usize {
|
||||
self.current_epoch
|
||||
}
|
||||
|
||||
/// Check if in warmup phase
|
||||
pub fn is_warmup(&self) -> bool {
|
||||
self.current_epoch < self.config.warmup_epochs
|
||||
}
|
||||
|
||||
/// Get configuration
|
||||
pub fn config(&self) -> &DARTSConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Get cells
|
||||
pub fn cells(&self) -> &[DARTSCell] {
|
||||
&self.cells
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_darts_config_new() {
|
||||
let config = DARTSConfig::new(0.001, 0.01, 50, 15);
|
||||
assert_eq!(config.learning_rate_arch, 0.001);
|
||||
assert_eq!(config.learning_rate_weights, 0.01);
|
||||
assert_eq!(config.num_epochs, 50);
|
||||
assert_eq!(config.warmup_epochs, 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_config_default() {
|
||||
let config = DARTSConfig::default();
|
||||
assert!(config.learning_rate_arch > 0.0);
|
||||
assert!(config.learning_rate_weights > 0.0);
|
||||
assert!(config.num_epochs > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_config_validate() {
|
||||
let config = DARTSConfig::default();
|
||||
assert!(config.validate().is_ok());
|
||||
|
||||
let invalid = DARTSConfig::new(-0.001, 0.01, 50, 15);
|
||||
assert!(invalid.validate().is_err());
|
||||
|
||||
let invalid = DARTSConfig::new(0.001, 0.01, 50, 60);
|
||||
assert!(invalid.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_config_with_temperature() {
|
||||
let config = DARTSConfig::default().with_temperature(2.0);
|
||||
assert_eq!(config.temperature, 2.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_softmax() {
|
||||
let values = vec![1.0, 2.0, 3.0];
|
||||
let probs = softmax(&values, 1.0);
|
||||
|
||||
// Check sum is 1.0
|
||||
let sum: f32 = probs.iter().sum();
|
||||
assert!((sum - 1.0).abs() < 1e-6);
|
||||
|
||||
// Check probabilities are positive
|
||||
assert!(probs.iter().all(|&p| p > 0.0));
|
||||
|
||||
// Check last value has highest probability
|
||||
assert!(probs[2] > probs[1]);
|
||||
assert!(probs[1] > probs[0]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_softmax_temperature() {
|
||||
let values = vec![1.0, 2.0, 3.0];
|
||||
|
||||
// Higher temperature makes distribution more uniform
|
||||
let probs_high_temp = softmax(&values, 10.0);
|
||||
let probs_low_temp = softmax(&values, 0.1);
|
||||
|
||||
// At high temperature, probabilities should be closer together
|
||||
let diff_high = probs_high_temp[2] - probs_high_temp[0];
|
||||
let diff_low = probs_low_temp[2] - probs_low_temp[0];
|
||||
|
||||
assert!(diff_high < diff_low);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mixed_op_new() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let edge = Edge::new(0, 2);
|
||||
let mixed_op = MixedOp::new(edge, 16, 32, 1, &device);
|
||||
|
||||
assert!(mixed_op.is_ok());
|
||||
let mixed_op = mixed_op.unwrap();
|
||||
assert_eq!(mixed_op.edge(), edge);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mixed_op_forward() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let edge = Edge::new(0, 2);
|
||||
let mixed_op = MixedOp::new(edge, 16, 32, 1, &device).unwrap();
|
||||
|
||||
// Create input
|
||||
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
||||
|
||||
// Create uniform weights
|
||||
let num_ops = OperationType::count();
|
||||
let weights = vec![1.0 / num_ops as f32; num_ops];
|
||||
|
||||
let result = mixed_op.forward(&input, &weights);
|
||||
|
||||
// Conv2d may not be fully implemented yet, so we just verify it doesn't panic
|
||||
if let Ok(output) = result {
|
||||
assert_eq!(output.shape(), &[2, 32, 8, 8]);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_cell_new() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let config = CellConfig::default_darts();
|
||||
let cell = DARTSCell::new(config, &device);
|
||||
|
||||
assert!(cell.is_ok());
|
||||
let cell = cell.unwrap();
|
||||
|
||||
// Check alphas are initialized
|
||||
let alphas = cell.alphas();
|
||||
assert!(!alphas.is_empty());
|
||||
|
||||
// Each edge should have alphas for all operations
|
||||
for alpha_vec in alphas.values() {
|
||||
assert_eq!(alpha_vec.len(), OperationType::count());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_cell_get_probabilities() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let config = CellConfig::default_darts();
|
||||
let cell = DARTSCell::new(config, &device).unwrap();
|
||||
|
||||
let probs = cell.get_probabilities(1.0);
|
||||
|
||||
// Check all probabilities sum to 1.0
|
||||
for prob_vec in probs.values() {
|
||||
let sum: f32 = prob_vec.iter().sum();
|
||||
assert!((sum - 1.0).abs() < 1e-5);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_cell_derive_architecture() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let config = CellConfig::default_darts();
|
||||
let cell = DARTSCell::new(config, &device).unwrap();
|
||||
|
||||
let derived = cell.derive_architecture();
|
||||
assert!(derived.is_ok());
|
||||
|
||||
let derived = derived.unwrap();
|
||||
assert!(derived.is_complete());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_new() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let config = DARTSConfig::default();
|
||||
let cell_configs = vec![CellConfig::default_darts()];
|
||||
|
||||
let darts = DARTS::new(config, cell_configs, &device);
|
||||
assert!(darts.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_default() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let darts = DARTS::default(&device);
|
||||
assert!(darts.is_ok());
|
||||
|
||||
let darts = darts.unwrap();
|
||||
assert_eq!(darts.cells().len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_step_warmup() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let mut darts = DARTS::default(&device).unwrap();
|
||||
|
||||
assert!(darts.is_warmup());
|
||||
|
||||
// During warmup, step should not update alphas significantly
|
||||
let alphas_before = darts.get_alphas()[0].clone();
|
||||
|
||||
let result = darts.step(0.5, 0.6, None, None);
|
||||
assert!(result.is_ok());
|
||||
|
||||
let alphas_after = darts.get_alphas()[0].clone();
|
||||
|
||||
// Alphas should remain similar during warmup
|
||||
for (edge, alpha_vec_after) in alphas_after {
|
||||
if let Some(alpha_vec_before) = alphas_before.get(&edge) {
|
||||
for (a, b) in alpha_vec_before.iter().zip(alpha_vec_after.iter()) {
|
||||
assert!((a - b).abs() < 1e-6);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_step_after_warmup() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let mut config = DARTSConfig::default();
|
||||
config.warmup_epochs = 0; // No warmup
|
||||
|
||||
let cell_configs = vec![CellConfig::default_darts()];
|
||||
let mut darts = DARTS::new(config, cell_configs, &device).unwrap();
|
||||
|
||||
assert!(!darts.is_warmup());
|
||||
|
||||
let result = darts.step(0.5, 0.6, None, None);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_derive_architecture() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let darts = DARTS::default(&device).unwrap();
|
||||
|
||||
let arch = darts.derive_architecture();
|
||||
assert!(arch.is_ok());
|
||||
|
||||
let arch = arch.unwrap();
|
||||
assert!(arch.validate().is_ok());
|
||||
assert_eq!(arch.num_cells(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_get_alphas() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let darts = DARTS::default(&device).unwrap();
|
||||
|
||||
let alphas = darts.get_alphas();
|
||||
assert_eq!(alphas.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_get_probabilities() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let darts = DARTS::default(&device).unwrap();
|
||||
|
||||
let probs = darts.get_probabilities();
|
||||
assert_eq!(probs.len(), 2);
|
||||
|
||||
// Check all probabilities are valid
|
||||
for prob_map in probs {
|
||||
for prob_vec in prob_map.values() {
|
||||
let sum: f32 = prob_vec.iter().sum();
|
||||
assert!((sum - 1.0).abs() < 1e-5);
|
||||
assert!(prob_vec.iter().all(|&p| p >= 0.0 && p <= 1.0));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_darts_current_epoch() {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let mut darts = DARTS::default(&device).unwrap();
|
||||
|
||||
assert_eq!(darts.current_epoch(), 0);
|
||||
|
||||
darts.step(0.5, 0.6, None, None).unwrap();
|
||||
assert_eq!(darts.current_epoch(), 1);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user