Skip to content

Commit 74e7a7c

Browse files
Michall00misadowskgreenrazer
authored
candle-onnx: Implement Selu operator (huggingface#2978)
* feat: added Elu operator * feat: added selu implementation in candle-nn * feat: added selu onnx operator implementation * test: added test for selu onnx operator * test: added more tests for selu * deleted elu * tests: added test based on onnx specification * lint --------- Co-authored-by: misadowsk <michalsad.protondynamic@gmail.com> Co-authored-by: keighbee <kb@huggingface.co>
1 parent 4c9d262 commit 74e7a7c

3 files changed

Lines changed: 216 additions & 0 deletions

File tree

candle-nn/src/ops.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,15 @@ pub fn leaky_relu(xs: &Tensor, negative_slope: f64) -> Result<Tensor> {
249249
xs.maximum(&zeros)? + xs.minimum(&zeros)? * negative_slope
250250
}
251251

252+
pub fn selu(xs: &Tensor, alpha: f32, gamma: f32) -> Result<Tensor> {
253+
let is_pos = xs.gt(0f32)?;
254+
let alpha_t = Tensor::full(alpha, xs.dims(), xs.device())?;
255+
let neg = xs.exp()?.mul(&alpha_t)?.sub(&alpha_t)?;
256+
let selu = is_pos.where_cond(xs, &neg)?;
257+
let gamma_t = Tensor::full(gamma, xs.dims(), xs.device())?;
258+
selu.broadcast_mul(&gamma_t)
259+
}
260+
252261
pub fn dropout(xs: &Tensor, drop_p: f32) -> Result<Tensor> {
253262
// This implementation is inefficient as it stores the full mask for the backward pass.
254263
// Instead we could just store the seed and have a specialized kernel that would both

candle-onnx/src/eval.rs

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2097,6 +2097,19 @@ fn simple_eval_(
20972097
let output = input.sign()?;
20982098
values.insert(node.output[0].clone(), output);
20992099
}
2100+
// https://onnx.ai/onnx/operators/onnx__Selu.html
2101+
"Selu" => {
2102+
let input = get(&node.input[0])?;
2103+
let alpha = get_attr_opt::<f32>(node, "alpha")?
2104+
.copied()
2105+
.unwrap_or(1.6732632);
2106+
let gamma = get_attr_opt::<f32>(node, "gamma")?
2107+
.copied()
2108+
.unwrap_or(1.050701);
2109+
let out = candle_nn::ops::selu(input, alpha as f32, gamma as f32)?;
2110+
values.insert(node.output[0].clone(), out);
2111+
}
2112+
21002113
// https://onnx.ai/onnx/operators/onnx__OneHot.html
21012114
"OneHot" => {
21022115
let indices = get(&node.input[0])?;

candle-onnx/tests/ops.rs

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6248,6 +6248,200 @@ fn test_sign_operation() -> Result<()> {
62486248
}
62496249

62506250
#[test]
6251+
fn test_selu_operator() -> Result<()> {
6252+
{
6253+
// Test 1: Default alpha and gamma
6254+
let default_graph = create_model_proto_with_graph(Some(GraphProto {
6255+
node: vec![NodeProto {
6256+
op_type: "Selu".to_string(),
6257+
domain: "".to_string(),
6258+
input: vec!["input".to_string()],
6259+
output: vec!["output".to_string()],
6260+
..Default::default()
6261+
}],
6262+
input: vec![ValueInfoProto {
6263+
name: "input".to_string(),
6264+
..Default::default()
6265+
}],
6266+
output: vec![ValueInfoProto {
6267+
name: "output".to_string(),
6268+
r#type: None,
6269+
..Default::default()
6270+
}],
6271+
..Default::default()
6272+
}));
6273+
6274+
let input = Tensor::from_vec(vec![-1.0f32, 0.0, 1.0, 2.0], (2, 2), &Device::Cpu)?;
6275+
let mut inputs = HashMap::new();
6276+
inputs.insert("input".to_string(), input);
6277+
6278+
let eval = simple_eval(&default_graph, inputs)?;
6279+
let output = eval.get("output").unwrap();
6280+
let out_vec = to_vec2_round(output, 4)?;
6281+
assert_eq!(out_vec, vec![vec![-1.1113, 0.0], vec![1.0507, 2.1014]]);
6282+
}
6283+
6284+
{
6285+
// Test 2: Change alpha and gamma
6286+
let custom_graph = create_model_proto_with_graph(Some(GraphProto {
6287+
node: vec![NodeProto {
6288+
op_type: "Selu".to_string(),
6289+
attribute: vec![
6290+
AttributeProto {
6291+
name: "alpha".to_string(),
6292+
r#type: AttributeType::Float as i32,
6293+
f: 2.0,
6294+
..Default::default()
6295+
},
6296+
AttributeProto {
6297+
name: "gamma".to_string(),
6298+
r#type: AttributeType::Float as i32,
6299+
f: 0.5,
6300+
..Default::default()
6301+
},
6302+
],
6303+
input: vec!["input".to_string()],
6304+
output: vec!["output".to_string()],
6305+
..Default::default()
6306+
}],
6307+
input: vec![ValueInfoProto {
6308+
name: "input".to_string(),
6309+
..Default::default()
6310+
}],
6311+
output: vec![ValueInfoProto {
6312+
name: "output".to_string(),
6313+
..Default::default()
6314+
}],
6315+
..Default::default()
6316+
}));
6317+
6318+
let input = Tensor::from_vec(vec![-1.0f32, 0.0, 1.0, 2.0], (2, 2), &Device::Cpu)?;
6319+
let mut inputs = HashMap::new();
6320+
inputs.insert("input".to_string(), input);
6321+
let eval = simple_eval(&custom_graph, inputs)?;
6322+
let output = eval.get("output").unwrap();
6323+
let out_vec = to_vec2_round(output, 4)?;
6324+
assert_eq!(out_vec, vec![vec![-0.6321, 0.0], vec![0.5, 1.0]]);
6325+
}
6326+
6327+
{
6328+
// Test 3: Different input values
6329+
let manual_graph = create_model_proto_with_graph(Some(GraphProto {
6330+
node: vec![NodeProto {
6331+
op_type: "Selu".to_string(),
6332+
domain: "".to_string(),
6333+
input: vec!["input".to_string()],
6334+
output: vec!["output".to_string()],
6335+
..Default::default()
6336+
}],
6337+
input: vec![ValueInfoProto {
6338+
name: "input".to_string(),
6339+
..Default::default()
6340+
}],
6341+
output: vec![ValueInfoProto {
6342+
name: "output".to_string(),
6343+
..Default::default()
6344+
}],
6345+
..Default::default()
6346+
}));
6347+
6348+
let expected = vec![-1.758, -1.7463, 0.0, 10.507];
6349+
6350+
let input = Tensor::from_vec(vec![-10.0f32, -5.0, 0.0, 10.0], (2, 2), &Device::Cpu)?;
6351+
let mut inputs = HashMap::new();
6352+
inputs.insert("input".to_string(), input);
6353+
let eval = simple_eval(&manual_graph, inputs)?;
6354+
let output = eval.get("output").unwrap();
6355+
let out_vec = to_vec2_round(output, 4)?;
6356+
assert_eq!(
6357+
out_vec,
6358+
vec![
6359+
vec![expected[0], expected[1]],
6360+
vec![expected[2], expected[3]]
6361+
]
6362+
);
6363+
}
6364+
6365+
{
6366+
// Test 4: Test based on https://github.com/onnx/onnx/blob/main/docs/Operators.md#Selu
6367+
let graph = create_model_proto_with_graph(Some(GraphProto {
6368+
node: vec![NodeProto {
6369+
op_type: "Selu".to_string(),
6370+
input: vec!["input".to_string()],
6371+
output: vec!["output".to_string()],
6372+
attribute: vec![
6373+
AttributeProto {
6374+
name: "alpha".to_string(),
6375+
r#type: AttributeType::Float as i32,
6376+
f: 2.0,
6377+
..Default::default()
6378+
},
6379+
AttributeProto {
6380+
name: "gamma".to_string(),
6381+
r#type: AttributeType::Float as i32,
6382+
f: 3.0,
6383+
..Default::default()
6384+
},
6385+
],
6386+
..Default::default()
6387+
}],
6388+
input: vec![ValueInfoProto {
6389+
name: "input".to_string(),
6390+
..Default::default()
6391+
}],
6392+
output: vec![ValueInfoProto {
6393+
name: "output".to_string(),
6394+
..Default::default()
6395+
}],
6396+
..Default::default()
6397+
}));
6398+
6399+
let input = Tensor::from_vec(vec![-1.0f32, 0.0, 1.0], (3,), &Device::Cpu)?;
6400+
let mut inputs = HashMap::new();
6401+
inputs.insert("input".to_string(), input);
6402+
6403+
let eval = simple_eval(&graph, inputs)?;
6404+
let output = eval.get("output").unwrap();
6405+
let out_vec = output.to_vec1::<f32>()?;
6406+
let expected = vec![-3.7927232, 0.0, 3.0];
6407+
6408+
for (o, e) in out_vec.iter().zip(expected.iter()) {
6409+
assert!((o - e).abs() < 1e-5, "Got {o}, expected {e}");
6410+
}
6411+
}
6412+
6413+
{
6414+
// Test 5: Empty tensor
6415+
let manual_graph = create_model_proto_with_graph(Some(GraphProto {
6416+
node: vec![NodeProto {
6417+
op_type: "Selu".to_string(),
6418+
domain: "".to_string(),
6419+
input: vec!["input".to_string()],
6420+
output: vec!["output".to_string()],
6421+
..Default::default()
6422+
}],
6423+
input: vec![ValueInfoProto {
6424+
name: "input".to_string(),
6425+
..Default::default()
6426+
}],
6427+
output: vec![ValueInfoProto {
6428+
name: "output".to_string(),
6429+
..Default::default()
6430+
}],
6431+
..Default::default()
6432+
}));
6433+
6434+
let input = Tensor::from_vec(vec![] as Vec<f32>, (0, 2), &Device::Cpu)?;
6435+
let mut inputs = HashMap::new();
6436+
inputs.insert("input".to_string(), input);
6437+
let eval = simple_eval(&manual_graph, inputs)?;
6438+
let output = eval.get("output").unwrap();
6439+
assert_eq!(output.dims(), &[0, 2]);
6440+
}
6441+
6442+
Ok(())
6443+
}
6444+
62516445
fn test_hard_swish() -> candle::Result<()> {
62526446
{
62536447
let manual_graph = create_model_proto_with_graph(Some(GraphProto {

0 commit comments

Comments
 (0)