Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,583 @@
|
||||
//! Node representation for ONNX operations.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::onnx_proto::{AttributeProto, NodeProto};
|
||||
|
||||
/// ONNX operator kinds.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub enum NodeKind {
|
||||
// Arithmetic
|
||||
/// Element-wise addition of two tensors.
|
||||
Add,
|
||||
/// Element-wise subtraction of two tensors.
|
||||
Sub,
|
||||
/// Element-wise multiplication of two tensors.
|
||||
Mul,
|
||||
/// Element-wise division of two tensors.
|
||||
Div,
|
||||
/// Element-wise negation of a tensor.
|
||||
Neg,
|
||||
/// Element-wise absolute value of a tensor.
|
||||
Abs,
|
||||
/// Element-wise square root of a tensor.
|
||||
Sqrt,
|
||||
/// Element-wise power operation.
|
||||
Pow,
|
||||
/// Element-wise exponential function.
|
||||
Exp,
|
||||
/// Element-wise natural logarithm.
|
||||
Log,
|
||||
/// Element-wise floor rounding.
|
||||
Floor,
|
||||
/// Element-wise ceiling rounding.
|
||||
Ceil,
|
||||
/// Element-wise rounding to nearest integer.
|
||||
Round,
|
||||
|
||||
// Matrix operations
|
||||
/// Matrix multiplication of two tensors.
|
||||
MatMul,
|
||||
/// General matrix multiplication with optional bias.
|
||||
Gemm,
|
||||
|
||||
// Activations
|
||||
/// Rectified Linear Unit activation function.
|
||||
Relu,
|
||||
/// Leaky Rectified Linear Unit activation function.
|
||||
LeakyRelu,
|
||||
/// Sigmoid activation function.
|
||||
Sigmoid,
|
||||
/// Hyperbolic tangent activation function.
|
||||
Tanh,
|
||||
/// Softmax activation function.
|
||||
Softmax,
|
||||
/// Log-softmax activation function.
|
||||
LogSoftmax,
|
||||
/// Gaussian Error Linear Unit activation function.
|
||||
Gelu,
|
||||
/// Sigmoid Linear Unit (Swish) activation function.
|
||||
Silu,
|
||||
/// Hard sigmoid activation function.
|
||||
HardSigmoid,
|
||||
/// Hard swish activation function.
|
||||
HardSwish,
|
||||
/// Exponential Linear Unit activation function.
|
||||
Elu,
|
||||
/// Scaled Exponential Linear Unit activation function.
|
||||
Selu,
|
||||
/// Continuously Differentiable Exponential Linear Unit.
|
||||
Celu,
|
||||
/// Softplus activation function.
|
||||
Softplus,
|
||||
/// Softsign activation function.
|
||||
Softsign,
|
||||
/// Mish activation function.
|
||||
Mish,
|
||||
|
||||
// Normalization
|
||||
/// Batch normalization operation.
|
||||
BatchNormalization,
|
||||
/// Layer normalization operation.
|
||||
LayerNormalization,
|
||||
/// Instance normalization operation.
|
||||
InstanceNormalization,
|
||||
/// Group normalization operation.
|
||||
GroupNormalization,
|
||||
/// Lp normalization operation.
|
||||
LpNormalization,
|
||||
|
||||
// Convolution
|
||||
/// Convolution operation.
|
||||
Conv,
|
||||
/// Transposed convolution (deconvolution) operation.
|
||||
ConvTranspose,
|
||||
|
||||
// Pooling
|
||||
/// Max pooling operation.
|
||||
MaxPool,
|
||||
/// Average pooling operation.
|
||||
AveragePool,
|
||||
/// Global average pooling over spatial dimensions.
|
||||
GlobalAveragePool,
|
||||
/// Global max pooling over spatial dimensions.
|
||||
GlobalMaxPool,
|
||||
/// Global Lp pooling over spatial dimensions.
|
||||
GlobalLpPool,
|
||||
|
||||
// Reduction
|
||||
/// Sum reduction along specified axes.
|
||||
ReduceSum,
|
||||
/// Mean reduction along specified axes.
|
||||
ReduceMean,
|
||||
/// Max reduction along specified axes.
|
||||
ReduceMax,
|
||||
/// Min reduction along specified axes.
|
||||
ReduceMin,
|
||||
/// Product reduction along specified axes.
|
||||
ReduceProd,
|
||||
/// L1 norm reduction along specified axes.
|
||||
ReduceL1,
|
||||
/// L2 norm reduction along specified axes.
|
||||
ReduceL2,
|
||||
/// Log of sum reduction along specified axes.
|
||||
ReduceLogSum,
|
||||
/// Log of sum of exponentials reduction along specified axes.
|
||||
ReduceLogSumExp,
|
||||
/// Sum of squares reduction along specified axes.
|
||||
ReduceSumSquare,
|
||||
|
||||
// Shape operations
|
||||
/// Reshape tensor to a new shape.
|
||||
Reshape,
|
||||
/// Transpose tensor dimensions.
|
||||
Transpose,
|
||||
/// Flatten tensor to 2D.
|
||||
Flatten,
|
||||
/// Remove dimensions of size 1.
|
||||
Squeeze,
|
||||
/// Insert dimensions of size 1.
|
||||
Unsqueeze,
|
||||
/// Concatenate tensors along an axis.
|
||||
Concat,
|
||||
/// Split tensor into multiple tensors.
|
||||
Split,
|
||||
/// Extract a slice from a tensor.
|
||||
Slice,
|
||||
/// Gather elements along an axis by index.
|
||||
Gather,
|
||||
/// Gather elements by index along specified axis.
|
||||
GatherElements,
|
||||
/// Gather slices from tensor using N-dimensional indices.
|
||||
GatherND,
|
||||
/// Scatter elements along an axis by index.
|
||||
Scatter,
|
||||
/// Scatter elements by index along specified axis.
|
||||
ScatterElements,
|
||||
/// Scatter slices into tensor using N-dimensional indices.
|
||||
ScatterND,
|
||||
/// Broadcast tensor to a larger shape.
|
||||
Expand,
|
||||
/// Tile tensor by repeating along dimensions.
|
||||
Tile,
|
||||
/// Pad tensor with constant or edge values.
|
||||
Pad,
|
||||
|
||||
// Comparison
|
||||
/// Element-wise equality comparison.
|
||||
Equal,
|
||||
/// Element-wise greater-than comparison.
|
||||
Greater,
|
||||
/// Element-wise greater-than-or-equal comparison.
|
||||
GreaterOrEqual,
|
||||
/// Element-wise less-than comparison.
|
||||
Less,
|
||||
/// Element-wise less-than-or-equal comparison.
|
||||
LessOrEqual,
|
||||
/// Element-wise logical NOT.
|
||||
Not,
|
||||
/// Element-wise logical AND.
|
||||
And,
|
||||
/// Element-wise logical OR.
|
||||
Or,
|
||||
/// Element-wise logical XOR.
|
||||
Xor,
|
||||
/// Conditional selection based on a mask tensor.
|
||||
Where,
|
||||
|
||||
// Type conversion
|
||||
/// Cast tensor to a different data type.
|
||||
Cast,
|
||||
/// Cast tensor to the same type as another tensor.
|
||||
CastLike,
|
||||
|
||||
// Constants
|
||||
/// Produce a constant tensor.
|
||||
Constant,
|
||||
/// Generate a tensor of a given shape filled with a constant.
|
||||
ConstantOfShape,
|
||||
/// Get the shape of a tensor as a 1D tensor.
|
||||
Shape,
|
||||
/// Get the total number of elements in a tensor.
|
||||
Size,
|
||||
|
||||
// RNN
|
||||
/// Long Short-Term Memory recurrent layer.
|
||||
LSTM,
|
||||
/// Gated Recurrent Unit recurrent layer.
|
||||
GRU,
|
||||
/// Simple recurrent neural network layer.
|
||||
RNN,
|
||||
|
||||
// Attention
|
||||
/// Single-head attention mechanism.
|
||||
Attention,
|
||||
/// Multi-head attention mechanism.
|
||||
MultiHeadAttention,
|
||||
|
||||
// Misc
|
||||
/// Dropout regularization layer.
|
||||
Dropout,
|
||||
/// Identity operation (pass-through).
|
||||
Identity,
|
||||
/// Clip values to a specified range.
|
||||
Clip,
|
||||
/// Gauss error function.
|
||||
Erf,
|
||||
/// Element-wise sine function.
|
||||
Sin,
|
||||
/// Element-wise cosine function.
|
||||
Cos,
|
||||
/// Element-wise tangent function.
|
||||
Tan,
|
||||
/// Element-wise arcsine function.
|
||||
Asin,
|
||||
/// Element-wise arccosine function.
|
||||
Acos,
|
||||
/// Element-wise arctangent function.
|
||||
Atan,
|
||||
/// Element-wise hyperbolic sine function.
|
||||
Sinh,
|
||||
/// Element-wise hyperbolic cosine function.
|
||||
Cosh,
|
||||
/// Element-wise inverse hyperbolic sine function.
|
||||
Asinh,
|
||||
/// Element-wise inverse hyperbolic cosine function.
|
||||
Acosh,
|
||||
/// Element-wise inverse hyperbolic tangent function.
|
||||
Atanh,
|
||||
/// Element-wise sign function.
|
||||
Sign,
|
||||
/// Element-wise reciprocal (1/x).
|
||||
Reciprocal,
|
||||
/// Element-wise minimum of input tensors.
|
||||
Min,
|
||||
/// Element-wise maximum of input tensors.
|
||||
Max,
|
||||
/// Element-wise mean of input tensors.
|
||||
Mean,
|
||||
/// Element-wise sum of input tensors.
|
||||
Sum,
|
||||
|
||||
// Embedding
|
||||
/// Embedding lookup operation.
|
||||
Embedding,
|
||||
/// One-hot encoding operation.
|
||||
OneHot,
|
||||
|
||||
// Resize
|
||||
/// Resize tensor using interpolation.
|
||||
Resize,
|
||||
/// Upsample tensor (deprecated, use Resize).
|
||||
Upsample,
|
||||
|
||||
// Custom/Unknown
|
||||
/// Custom or unsupported operator with its name.
|
||||
Custom(String),
|
||||
}
|
||||
|
||||
impl NodeKind {
|
||||
/// Parse from ONNX op_type string.
|
||||
pub fn from_op_type(op_type: &str) -> Self {
|
||||
match op_type {
|
||||
// Arithmetic
|
||||
"Add" => NodeKind::Add,
|
||||
"Sub" => NodeKind::Sub,
|
||||
"Mul" => NodeKind::Mul,
|
||||
"Div" => NodeKind::Div,
|
||||
"Neg" => NodeKind::Neg,
|
||||
"Abs" => NodeKind::Abs,
|
||||
"Sqrt" => NodeKind::Sqrt,
|
||||
"Pow" => NodeKind::Pow,
|
||||
"Exp" => NodeKind::Exp,
|
||||
"Log" => NodeKind::Log,
|
||||
"Floor" => NodeKind::Floor,
|
||||
"Ceil" => NodeKind::Ceil,
|
||||
"Round" => NodeKind::Round,
|
||||
|
||||
// Matrix
|
||||
"MatMul" => NodeKind::MatMul,
|
||||
"Gemm" => NodeKind::Gemm,
|
||||
|
||||
// Activations
|
||||
"Relu" => NodeKind::Relu,
|
||||
"LeakyRelu" => NodeKind::LeakyRelu,
|
||||
"Sigmoid" => NodeKind::Sigmoid,
|
||||
"Tanh" => NodeKind::Tanh,
|
||||
"Softmax" => NodeKind::Softmax,
|
||||
"LogSoftmax" => NodeKind::LogSoftmax,
|
||||
"Gelu" => NodeKind::Gelu,
|
||||
"Silu" => NodeKind::Silu,
|
||||
"HardSigmoid" => NodeKind::HardSigmoid,
|
||||
"HardSwish" => NodeKind::HardSwish,
|
||||
"Elu" => NodeKind::Elu,
|
||||
"Selu" => NodeKind::Selu,
|
||||
"Celu" => NodeKind::Celu,
|
||||
"Softplus" => NodeKind::Softplus,
|
||||
"Softsign" => NodeKind::Softsign,
|
||||
"Mish" => NodeKind::Mish,
|
||||
|
||||
// Normalization
|
||||
"BatchNormalization" => NodeKind::BatchNormalization,
|
||||
"LayerNormalization" => NodeKind::LayerNormalization,
|
||||
"InstanceNormalization" => NodeKind::InstanceNormalization,
|
||||
"GroupNormalization" => NodeKind::GroupNormalization,
|
||||
"LpNormalization" => NodeKind::LpNormalization,
|
||||
|
||||
// Convolution
|
||||
"Conv" => NodeKind::Conv,
|
||||
"ConvTranspose" => NodeKind::ConvTranspose,
|
||||
|
||||
// Pooling
|
||||
"MaxPool" => NodeKind::MaxPool,
|
||||
"AveragePool" => NodeKind::AveragePool,
|
||||
"GlobalAveragePool" => NodeKind::GlobalAveragePool,
|
||||
"GlobalMaxPool" => NodeKind::GlobalMaxPool,
|
||||
"GlobalLpPool" => NodeKind::GlobalLpPool,
|
||||
|
||||
// Reduction
|
||||
"ReduceSum" => NodeKind::ReduceSum,
|
||||
"ReduceMean" => NodeKind::ReduceMean,
|
||||
"ReduceMax" => NodeKind::ReduceMax,
|
||||
"ReduceMin" => NodeKind::ReduceMin,
|
||||
"ReduceProd" => NodeKind::ReduceProd,
|
||||
"ReduceL1" => NodeKind::ReduceL1,
|
||||
"ReduceL2" => NodeKind::ReduceL2,
|
||||
"ReduceLogSum" => NodeKind::ReduceLogSum,
|
||||
"ReduceLogSumExp" => NodeKind::ReduceLogSumExp,
|
||||
"ReduceSumSquare" => NodeKind::ReduceSumSquare,
|
||||
|
||||
// Shape
|
||||
"Reshape" => NodeKind::Reshape,
|
||||
"Transpose" => NodeKind::Transpose,
|
||||
"Flatten" => NodeKind::Flatten,
|
||||
"Squeeze" => NodeKind::Squeeze,
|
||||
"Unsqueeze" => NodeKind::Unsqueeze,
|
||||
"Concat" => NodeKind::Concat,
|
||||
"Split" => NodeKind::Split,
|
||||
"Slice" => NodeKind::Slice,
|
||||
"Gather" => NodeKind::Gather,
|
||||
"GatherElements" => NodeKind::GatherElements,
|
||||
"GatherND" => NodeKind::GatherND,
|
||||
"Scatter" => NodeKind::Scatter,
|
||||
"ScatterElements" => NodeKind::ScatterElements,
|
||||
"ScatterND" => NodeKind::ScatterND,
|
||||
"Expand" => NodeKind::Expand,
|
||||
"Tile" => NodeKind::Tile,
|
||||
"Pad" => NodeKind::Pad,
|
||||
|
||||
// Comparison
|
||||
"Equal" => NodeKind::Equal,
|
||||
"Greater" => NodeKind::Greater,
|
||||
"GreaterOrEqual" => NodeKind::GreaterOrEqual,
|
||||
"Less" => NodeKind::Less,
|
||||
"LessOrEqual" => NodeKind::LessOrEqual,
|
||||
"Not" => NodeKind::Not,
|
||||
"And" => NodeKind::And,
|
||||
"Or" => NodeKind::Or,
|
||||
"Xor" => NodeKind::Xor,
|
||||
"Where" => NodeKind::Where,
|
||||
|
||||
// Type
|
||||
"Cast" => NodeKind::Cast,
|
||||
"CastLike" => NodeKind::CastLike,
|
||||
|
||||
// Constants
|
||||
"Constant" => NodeKind::Constant,
|
||||
"ConstantOfShape" => NodeKind::ConstantOfShape,
|
||||
"Shape" => NodeKind::Shape,
|
||||
"Size" => NodeKind::Size,
|
||||
|
||||
// RNN
|
||||
"LSTM" => NodeKind::LSTM,
|
||||
"GRU" => NodeKind::GRU,
|
||||
"RNN" => NodeKind::RNN,
|
||||
|
||||
// Attention
|
||||
"Attention" => NodeKind::Attention,
|
||||
"MultiHeadAttention" => NodeKind::MultiHeadAttention,
|
||||
|
||||
// Misc
|
||||
"Dropout" => NodeKind::Dropout,
|
||||
"Identity" => NodeKind::Identity,
|
||||
"Clip" => NodeKind::Clip,
|
||||
"Erf" => NodeKind::Erf,
|
||||
"Sin" => NodeKind::Sin,
|
||||
"Cos" => NodeKind::Cos,
|
||||
"Tan" => NodeKind::Tan,
|
||||
"Asin" => NodeKind::Asin,
|
||||
"Acos" => NodeKind::Acos,
|
||||
"Atan" => NodeKind::Atan,
|
||||
"Sinh" => NodeKind::Sinh,
|
||||
"Cosh" => NodeKind::Cosh,
|
||||
"Asinh" => NodeKind::Asinh,
|
||||
"Acosh" => NodeKind::Acosh,
|
||||
"Atanh" => NodeKind::Atanh,
|
||||
"Sign" => NodeKind::Sign,
|
||||
"Reciprocal" => NodeKind::Reciprocal,
|
||||
"Min" => NodeKind::Min,
|
||||
"Max" => NodeKind::Max,
|
||||
"Mean" => NodeKind::Mean,
|
||||
"Sum" => NodeKind::Sum,
|
||||
|
||||
// Embedding
|
||||
"Embedding" => NodeKind::Embedding,
|
||||
"OneHot" => NodeKind::OneHot,
|
||||
|
||||
// Resize
|
||||
"Resize" => NodeKind::Resize,
|
||||
"Upsample" => NodeKind::Upsample,
|
||||
|
||||
// Unknown
|
||||
other => NodeKind::Custom(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if this is a supported operator.
|
||||
pub fn is_supported(&self) -> bool {
|
||||
!matches!(self, NodeKind::Custom(_))
|
||||
}
|
||||
}
|
||||
|
||||
/// Attribute value types.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum AttributeValue {
|
||||
/// Floating point value.
|
||||
Float(f32),
|
||||
/// Integer value.
|
||||
Int(i64),
|
||||
/// String value.
|
||||
String(String),
|
||||
/// List of floats.
|
||||
Floats(Vec<f32>),
|
||||
/// List of integers.
|
||||
Ints(Vec<i64>),
|
||||
/// List of strings.
|
||||
Strings(Vec<String>),
|
||||
}
|
||||
|
||||
/// A node in the ONNX graph.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Node {
|
||||
/// Node name (may be empty).
|
||||
pub name: String,
|
||||
/// Operator kind.
|
||||
pub kind: NodeKind,
|
||||
/// Input tensor names.
|
||||
pub inputs: Vec<String>,
|
||||
/// Output tensor names.
|
||||
pub outputs: Vec<String>,
|
||||
/// Node attributes.
|
||||
pub attributes: HashMap<String, AttributeValue>,
|
||||
}
|
||||
|
||||
impl Node {
|
||||
/// Create from ONNX NodeProto.
|
||||
pub fn from_proto(proto: &NodeProto) -> Result<Self> {
|
||||
let name = proto.name.clone();
|
||||
let kind = NodeKind::from_op_type(&proto.op_type);
|
||||
let inputs = proto.input.clone();
|
||||
let outputs = proto.output.clone();
|
||||
let attributes = Self::parse_attributes(&proto.attribute)?;
|
||||
|
||||
Ok(Self {
|
||||
name,
|
||||
kind,
|
||||
inputs,
|
||||
outputs,
|
||||
attributes,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_attributes(attrs: &[AttributeProto]) -> Result<HashMap<String, AttributeValue>> {
|
||||
let mut result = HashMap::new();
|
||||
|
||||
for attr in attrs {
|
||||
let value = Self::parse_attribute_value(attr)?;
|
||||
result.insert(attr.name.clone(), value);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
fn parse_attribute_value(attr: &AttributeProto) -> Result<AttributeValue> {
|
||||
// Check type field first, then fall back to presence of data
|
||||
match attr.r#type {
|
||||
1 => Ok(AttributeValue::Float(attr.f)),
|
||||
2 => Ok(AttributeValue::Int(attr.i)),
|
||||
3 => Ok(AttributeValue::String(
|
||||
String::from_utf8_lossy(&attr.s).to_string(),
|
||||
)),
|
||||
6 => Ok(AttributeValue::Floats(attr.floats.clone())),
|
||||
7 => Ok(AttributeValue::Ints(attr.ints.clone())),
|
||||
8 => Ok(AttributeValue::Strings(
|
||||
attr.strings
|
||||
.iter()
|
||||
.map(|s| String::from_utf8_lossy(s).to_string())
|
||||
.collect(),
|
||||
)),
|
||||
_ => {
|
||||
// Fall back to checking which field has data
|
||||
if attr.f != 0.0 {
|
||||
Ok(AttributeValue::Float(attr.f))
|
||||
} else if attr.i != 0 {
|
||||
Ok(AttributeValue::Int(attr.i))
|
||||
} else if !attr.s.is_empty() {
|
||||
Ok(AttributeValue::String(
|
||||
String::from_utf8_lossy(&attr.s).to_string(),
|
||||
))
|
||||
} else if !attr.floats.is_empty() {
|
||||
Ok(AttributeValue::Floats(attr.floats.clone()))
|
||||
} else if !attr.ints.is_empty() {
|
||||
Ok(AttributeValue::Ints(attr.ints.clone()))
|
||||
} else {
|
||||
// Default to int 0
|
||||
Ok(AttributeValue::Int(0))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get an integer attribute.
|
||||
pub fn get_int(&self, name: &str) -> Option<i64> {
|
||||
match self.attributes.get(name) {
|
||||
Some(AttributeValue::Int(v)) => Some(*v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a float attribute.
|
||||
pub fn get_float(&self, name: &str) -> Option<f32> {
|
||||
match self.attributes.get(name) {
|
||||
Some(AttributeValue::Float(v)) => Some(*v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a string attribute.
|
||||
pub fn get_string(&self, name: &str) -> Option<&str> {
|
||||
match self.attributes.get(name) {
|
||||
Some(AttributeValue::String(v)) => Some(v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get an integer list attribute.
|
||||
pub fn get_ints(&self, name: &str) -> Option<&[i64]> {
|
||||
match self.attributes.get(name) {
|
||||
Some(AttributeValue::Ints(v)) => Some(v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a float list attribute.
|
||||
pub fn get_floats(&self, name: &str) -> Option<&[f32]> {
|
||||
match self.attributes.get(name) {
|
||||
Some(AttributeValue::Floats(v)) => Some(v),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user