581 lines
19 KiB
Rust
581 lines
19 KiB
Rust
//! Operation primitives for Neural Architecture Search
|
|
|
|
use crate::error::{NASError, Result};
|
|
use rtx_nn::layers::Module;
|
|
use rtx_tensor::{Device, Tensor};
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Operation types available in the search space
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Hash)]
|
|
pub enum OperationType {
|
|
/// Identity/skip connection
|
|
Identity,
|
|
/// Zero/no connection
|
|
Zero,
|
|
/// 3x3 convolution
|
|
Conv3x3,
|
|
/// 5x5 convolution
|
|
Conv5x5,
|
|
/// Separable 3x3 convolution
|
|
SepConv3x3,
|
|
/// Separable 5x5 convolution
|
|
SepConv5x5,
|
|
/// Dilated 3x3 convolution
|
|
DilConv3x3,
|
|
/// 3x3 max pooling
|
|
MaxPool3x3,
|
|
/// 3x3 average pooling
|
|
AvgPool3x3,
|
|
}
|
|
|
|
impl OperationType {
|
|
/// Get all available operation types
|
|
pub fn all() -> Vec<Self> {
|
|
vec![
|
|
Self::Identity,
|
|
Self::Zero,
|
|
Self::Conv3x3,
|
|
Self::Conv5x5,
|
|
Self::SepConv3x3,
|
|
Self::SepConv5x5,
|
|
Self::DilConv3x3,
|
|
Self::MaxPool3x3,
|
|
Self::AvgPool3x3,
|
|
]
|
|
}
|
|
|
|
/// Get the number of available operations
|
|
pub fn count() -> usize {
|
|
Self::all().len()
|
|
}
|
|
|
|
/// Get operation name
|
|
pub fn name(&self) -> &'static str {
|
|
match self {
|
|
Self::Identity => "identity",
|
|
Self::Zero => "zero",
|
|
Self::Conv3x3 => "conv_3x3",
|
|
Self::Conv5x5 => "conv_5x5",
|
|
Self::SepConv3x3 => "sep_conv_3x3",
|
|
Self::SepConv5x5 => "sep_conv_5x5",
|
|
Self::DilConv3x3 => "dil_conv_3x3",
|
|
Self::MaxPool3x3 => "max_pool_3x3",
|
|
Self::AvgPool3x3 => "avg_pool_3x3",
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Configuration for an operation
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct OperationConfig {
|
|
/// Operation type
|
|
pub op_type: OperationType,
|
|
/// Number of input channels
|
|
pub in_channels: usize,
|
|
/// Number of output channels
|
|
pub out_channels: usize,
|
|
/// Stride
|
|
pub stride: usize,
|
|
/// Whether to apply affine transformation (for normalization layers)
|
|
pub affine: bool,
|
|
}
|
|
|
|
impl OperationConfig {
|
|
/// Create a new operation configuration
|
|
pub fn new(
|
|
op_type: OperationType,
|
|
in_channels: usize,
|
|
out_channels: usize,
|
|
stride: usize,
|
|
) -> Self {
|
|
Self {
|
|
op_type,
|
|
in_channels,
|
|
out_channels,
|
|
stride,
|
|
affine: true,
|
|
}
|
|
}
|
|
|
|
/// Validate the configuration
|
|
pub fn validate(&self) -> Result<()> {
|
|
if self.in_channels == 0 {
|
|
return Err(NASError::InvalidConfig(
|
|
"in_channels must be greater than 0".into(),
|
|
));
|
|
}
|
|
if self.out_channels == 0 {
|
|
return Err(NASError::InvalidConfig(
|
|
"out_channels must be greater than 0".into(),
|
|
));
|
|
}
|
|
if self.stride == 0 {
|
|
return Err(NASError::InvalidConfig(
|
|
"stride must be greater than 0".into(),
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// An executable operation in the search space
|
|
#[derive(Debug)]
|
|
pub struct Operation {
|
|
config: OperationConfig,
|
|
device: Device,
|
|
}
|
|
|
|
impl Operation {
|
|
/// Create a new operation
|
|
pub fn new(config: OperationConfig, device: &Device) -> Result<Self> {
|
|
config.validate()?;
|
|
Ok(Self {
|
|
config,
|
|
device: device.clone(),
|
|
})
|
|
}
|
|
|
|
/// Get the operation type
|
|
pub fn op_type(&self) -> OperationType {
|
|
self.config.op_type
|
|
}
|
|
|
|
/// Get the operation configuration
|
|
pub fn config(&self) -> &OperationConfig {
|
|
&self.config
|
|
}
|
|
|
|
/// Execute the operation (forward pass)
|
|
pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
|
match self.config.op_type {
|
|
OperationType::Identity => self.identity(input),
|
|
OperationType::Zero => self.zero(input),
|
|
OperationType::Conv3x3 => self.conv(input, 3),
|
|
OperationType::Conv5x5 => self.conv(input, 5),
|
|
OperationType::SepConv3x3 => self.sep_conv(input, 3),
|
|
OperationType::SepConv5x5 => self.sep_conv(input, 5),
|
|
OperationType::DilConv3x3 => self.dil_conv(input, 3),
|
|
OperationType::MaxPool3x3 => self.max_pool(input),
|
|
OperationType::AvgPool3x3 => self.avg_pool(input),
|
|
}
|
|
}
|
|
|
|
/// Identity operation (skip connection)
|
|
fn identity(&self, input: &Tensor) -> Result<Tensor> {
|
|
// If channels match and stride is 1, just return input
|
|
let input_shape = input.shape();
|
|
if input_shape.len() != 4 {
|
|
return Err(NASError::OperationError(
|
|
"Input must be 4D tensor [B, C, H, W]".into(),
|
|
));
|
|
}
|
|
|
|
let in_c = input_shape[1];
|
|
if in_c == self.config.out_channels && self.config.stride == 1 {
|
|
Ok(input.clone())
|
|
} else {
|
|
// Need to adjust channels or downsample
|
|
let mut result = input.clone();
|
|
|
|
// Handle stride (downsampling)
|
|
if self.config.stride > 1 {
|
|
result = self.downsample(&result)?;
|
|
}
|
|
|
|
// Handle channel mismatch
|
|
if in_c != self.config.out_channels {
|
|
result = self.adjust_channels(&result)?;
|
|
}
|
|
|
|
Ok(result)
|
|
}
|
|
}
|
|
|
|
/// Zero operation (returns zeros)
|
|
fn zero(&self, input: &Tensor) -> Result<Tensor> {
|
|
let input_shape = input.shape();
|
|
if input_shape.len() != 4 {
|
|
return Err(NASError::OperationError(
|
|
"Input must be 4D tensor [B, C, H, W]".into(),
|
|
));
|
|
}
|
|
|
|
let b = input_shape[0];
|
|
let h = input_shape[2] / self.config.stride;
|
|
let w = input_shape[3] / self.config.stride;
|
|
|
|
let output_shape = vec![b, self.config.out_channels, h, w];
|
|
Ok(Tensor::zeros(output_shape, &self.device)?)
|
|
}
|
|
|
|
/// Standard convolution
|
|
fn conv(&self, input: &Tensor, kernel_size: usize) -> Result<Tensor> {
|
|
use rtx_nn::layers::conv::{Conv2d, Conv2dConfig};
|
|
|
|
let padding = kernel_size / 2;
|
|
let mut config = Conv2dConfig::square(
|
|
self.config.in_channels,
|
|
self.config.out_channels,
|
|
kernel_size,
|
|
);
|
|
config.stride = (self.config.stride, self.config.stride);
|
|
config.padding = (padding, padding);
|
|
|
|
let conv = Conv2d::from_config(config, &self.device)?;
|
|
conv.forward(input).map_err(Into::into)
|
|
}
|
|
|
|
/// Separable convolution (depthwise + pointwise)
|
|
fn sep_conv(&self, input: &Tensor, kernel_size: usize) -> Result<Tensor> {
|
|
use rtx_nn::layers::conv::{Conv2d, Conv2dConfig};
|
|
|
|
let padding = kernel_size / 2;
|
|
|
|
// Depthwise convolution
|
|
let mut depthwise_config = Conv2dConfig::square(
|
|
self.config.in_channels,
|
|
self.config.in_channels,
|
|
kernel_size,
|
|
);
|
|
depthwise_config.stride = (self.config.stride, self.config.stride);
|
|
depthwise_config.padding = (padding, padding);
|
|
depthwise_config.groups = self.config.in_channels;
|
|
|
|
let depthwise = Conv2d::from_config(depthwise_config, &self.device)?;
|
|
let intermediate = depthwise.forward(input)?;
|
|
|
|
// Pointwise convolution
|
|
let pointwise_config =
|
|
Conv2dConfig::square(self.config.in_channels, self.config.out_channels, 1);
|
|
let pointwise = Conv2d::from_config(pointwise_config, &self.device)?;
|
|
|
|
pointwise.forward(&intermediate).map_err(Into::into)
|
|
}
|
|
|
|
/// Dilated convolution
|
|
fn dil_conv(&self, input: &Tensor, kernel_size: usize) -> Result<Tensor> {
|
|
use rtx_nn::layers::conv::{Conv2d, Conv2dConfig};
|
|
|
|
let dilation = 2;
|
|
let padding = dilation * (kernel_size / 2);
|
|
|
|
let mut config = Conv2dConfig::square(
|
|
self.config.in_channels,
|
|
self.config.out_channels,
|
|
kernel_size,
|
|
);
|
|
config.stride = (self.config.stride, self.config.stride);
|
|
config.padding = (padding, padding);
|
|
config.dilation = (dilation, dilation);
|
|
|
|
let conv = Conv2d::from_config(config, &self.device)?;
|
|
conv.forward(input).map_err(Into::into)
|
|
}
|
|
|
|
/// Max pooling
|
|
fn max_pool(&self, input: &Tensor) -> Result<Tensor> {
|
|
use rtx_nn::layers::pooling::{MaxPool2d, MaxPool2dConfig};
|
|
|
|
let mut config = MaxPool2dConfig::new(3);
|
|
config.stride = Some((self.config.stride, self.config.stride));
|
|
config.padding = (1, 1);
|
|
|
|
let pool = MaxPool2d::from_config(config, &self.device)?;
|
|
let pooled = pool.forward(input)?;
|
|
|
|
// Adjust channels if needed
|
|
if input.shape()[1] == self.config.out_channels {
|
|
Ok(pooled)
|
|
} else {
|
|
self.adjust_channels(&pooled)
|
|
}
|
|
}
|
|
|
|
/// Average pooling
|
|
fn avg_pool(&self, input: &Tensor) -> Result<Tensor> {
|
|
use rtx_nn::layers::pooling::{AvgPool2d, AvgPool2dConfig};
|
|
|
|
let mut config = AvgPool2dConfig::new(3);
|
|
config.stride = Some((self.config.stride, self.config.stride));
|
|
config.padding = (1, 1);
|
|
|
|
let pool = AvgPool2d::from_config(config, &self.device)?;
|
|
let pooled = pool.forward(input)?;
|
|
|
|
// Adjust channels if needed
|
|
if input.shape()[1] == self.config.out_channels {
|
|
Ok(pooled)
|
|
} else {
|
|
self.adjust_channels(&pooled)
|
|
}
|
|
}
|
|
|
|
/// Downsample spatially
|
|
fn downsample(&self, input: &Tensor) -> Result<Tensor> {
|
|
use rtx_nn::layers::pooling::{AvgPool2d, AvgPool2dConfig};
|
|
|
|
let mut config = AvgPool2dConfig::new(1);
|
|
config.stride = Some((self.config.stride, self.config.stride));
|
|
|
|
let pool = AvgPool2d::from_config(config, &self.device)?;
|
|
pool.forward(input).map_err(Into::into)
|
|
}
|
|
|
|
/// Adjust number of channels
|
|
fn adjust_channels(&self, input: &Tensor) -> Result<Tensor> {
|
|
use rtx_nn::layers::conv::Conv2d;
|
|
|
|
let in_channels = input.shape()[1];
|
|
let conv = Conv2d::new(in_channels, self.config.out_channels, 1, &self.device)?;
|
|
|
|
conv.forward(input).map_err(Into::into)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_operation_type_all() {
|
|
let all_ops = OperationType::all();
|
|
assert_eq!(all_ops.len(), 9);
|
|
assert!(all_ops.contains(&OperationType::Identity));
|
|
assert!(all_ops.contains(&OperationType::Zero));
|
|
assert!(all_ops.contains(&OperationType::Conv3x3));
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_type_count() {
|
|
assert_eq!(OperationType::count(), 9);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_type_name() {
|
|
assert_eq!(OperationType::Identity.name(), "identity");
|
|
assert_eq!(OperationType::Conv3x3.name(), "conv_3x3");
|
|
assert_eq!(OperationType::MaxPool3x3.name(), "max_pool_3x3");
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_config_new() {
|
|
let config = OperationConfig::new(OperationType::Conv3x3, 16, 32, 1);
|
|
assert_eq!(config.op_type, OperationType::Conv3x3);
|
|
assert_eq!(config.in_channels, 16);
|
|
assert_eq!(config.out_channels, 32);
|
|
assert_eq!(config.stride, 1);
|
|
assert!(config.affine);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_config_validate() {
|
|
let config = OperationConfig::new(OperationType::Conv3x3, 16, 32, 1);
|
|
assert!(config.validate().is_ok());
|
|
|
|
let invalid_config = OperationConfig::new(OperationType::Conv3x3, 0, 32, 1);
|
|
assert!(invalid_config.validate().is_err());
|
|
|
|
let invalid_config = OperationConfig::new(OperationType::Conv3x3, 16, 0, 1);
|
|
assert!(invalid_config.validate().is_err());
|
|
|
|
let invalid_config = OperationConfig::new(OperationType::Conv3x3, 16, 32, 0);
|
|
assert!(invalid_config.validate().is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_creation() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Identity, 16, 16, 1);
|
|
let op = Operation::new(config, &device);
|
|
assert!(op.is_ok());
|
|
|
|
let op = op.unwrap();
|
|
assert_eq!(op.op_type(), OperationType::Identity);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_zero() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Zero, 16, 32, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let output = op.forward(&input).unwrap();
|
|
|
|
// Output should be [2, 32, 8, 8] with all zeros
|
|
assert_eq!(output.shape(), &[2, 32, 8, 8]);
|
|
|
|
// Verify it's zeros (check sum is close to 0)
|
|
let sum = output.sum(None).unwrap();
|
|
let sum_scalar: f32 = sum.item().unwrap();
|
|
assert!(sum_scalar.abs() < 1e-6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_identity_same_channels() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Identity, 16, 16, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let output = op.forward(&input).unwrap();
|
|
|
|
// Output should have same shape
|
|
assert_eq!(output.shape(), input.shape());
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_identity_different_channels() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Identity, 16, 32, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let result = op.forward(&input);
|
|
|
|
// Conv2d may not be fully implemented yet (used for channel adjustment)
|
|
if let Ok(output) = result {
|
|
assert_eq!(output.shape(), &[2, 32, 8, 8]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_identity_with_stride() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Identity, 16, 16, 2);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let output = op.forward(&input).unwrap();
|
|
|
|
// Output should be downsampled
|
|
assert_eq!(output.shape(), &[2, 16, 4, 4]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_conv3x3() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Conv3x3, 16, 32, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let result = op.forward(&input);
|
|
|
|
// Conv2d may not be fully implemented yet, so we just check it doesn't panic
|
|
// and produces some output with correct shape if successful
|
|
if let Ok(output) = result {
|
|
assert_eq!(output.shape(), &[2, 32, 8, 8]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_conv5x5() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Conv5x5, 16, 32, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let result = op.forward(&input);
|
|
|
|
// Conv2d may not be fully implemented yet
|
|
if let Ok(output) = result {
|
|
assert_eq!(output.shape(), &[2, 32, 8, 8]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_separable_conv() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::SepConv3x3, 16, 32, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let result = op.forward(&input);
|
|
|
|
// Conv2d may not be fully implemented yet
|
|
if let Ok(output) = result {
|
|
assert_eq!(output.shape(), &[2, 32, 8, 8]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_dilated_conv() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::DilConv3x3, 16, 32, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let result = op.forward(&input);
|
|
|
|
// Conv2d may not be fully implemented yet
|
|
if let Ok(output) = result {
|
|
assert_eq!(output.shape(), &[2, 32, 8, 8]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_max_pool() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::MaxPool3x3, 16, 16, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let output = op.forward(&input).unwrap();
|
|
|
|
// Output should have correct shape
|
|
assert_eq!(output.shape(), &[2, 16, 8, 8]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_avg_pool() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::AvgPool3x3, 16, 16, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let output = op.forward(&input).unwrap();
|
|
|
|
// Output should have correct shape
|
|
assert_eq!(output.shape(), &[2, 16, 8, 8]);
|
|
}
|
|
|
|
#[test]
|
|
fn test_operation_stride() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Conv3x3, 16, 32, 2);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create input tensor [B=2, C=16, H=8, W=8]
|
|
let input = Tensor::randn(&[2, 16, 8, 8], &device).unwrap();
|
|
let result = op.forward(&input);
|
|
|
|
// Conv2d may not be fully implemented yet
|
|
if let Ok(output) = result {
|
|
assert_eq!(output.shape(), &[2, 32, 4, 4]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_invalid_input_shape() {
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
let config = OperationConfig::new(OperationType::Identity, 16, 16, 1);
|
|
let op = Operation::new(config, &device).unwrap();
|
|
|
|
// Create 3D input tensor (invalid)
|
|
let input = Tensor::randn(&[2, 16, 8], &device).unwrap();
|
|
let result = op.forward(&input);
|
|
|
|
assert!(result.is_err());
|
|
}
|
|
}
|