584 lines
17 KiB
Rust
584 lines
17 KiB
Rust
//! 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,
|
|
}
|
|
}
|
|
}
|