From 2648e797c22ea9a25dd43cf6ecf0222dfd221170 Mon Sep 17 00:00:00 2001 From: Laurent Mazare Date: Tue, 5 Dec 2023 07:09:31 +0100 Subject: [PATCH] Use the proper broadcasting for prelu. (#1406) --- candle-nn/src/activation.rs | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/candle-nn/src/activation.rs b/candle-nn/src/activation.rs index 8b9a8785..c7dd359f 100644 --- a/candle-nn/src/activation.rs +++ b/candle-nn/src/activation.rs @@ -65,6 +65,15 @@ impl candle::Module for PReLU { fn forward(&self, xs: &Tensor) -> Result { let weight = if self.is_scalar { self.weight.reshape(())? + } else if xs.rank() >= 2 { + let num_channels = xs.dim(1)?; + let num_weights = self.weight.elem_count(); + if num_weights != num_channels { + candle::bail!("error in prelu: unexpected number of channels for the input, got {num_channels}, weight dim is {num_weights}") + } + let mut s = vec![1; xs.rank()]; + s[1] = self.weight.elem_count(); + self.weight.broadcast_as(s)? } else { self.weight.clone() }; @@ -78,11 +87,13 @@ impl candle::Module for PReLU { /// This uses some default name for weights, namely `"weight"`. /// # Arguments /// -/// * `num_parameters` - The number of parameters. Use `None` to have as single trainable value -/// and `Some` for a 1D vector with the appropriate number of features. -pub fn prelu(num_parameters: Option, vs: crate::VarBuilder) -> Result { +/// * `num_channels` - The number of channels. Use `None` to have as single trainable value and +/// `Some` for a 1D vector with the appropriate number of channels. When applying the `forward` +/// function, the input tensor shape `s` should either be one dimension with this number of +/// channels or if `s.len() >= 2` it should have `s[1]` equal to this number. +pub fn prelu(num_channels: Option, vs: crate::VarBuilder) -> Result { let init_ws = crate::init::Init::Const(0.25); // When using a scalar weight, the PyTorch encoding is to use a 1d vector of length 1. - let ws = vs.get_with_hints((num_parameters.unwrap_or(1),), "weight", init_ws)?; - Ok(PReLU::new(ws, num_parameters.is_none())) + let ws = vs.get_with_hints((num_channels.unwrap_or(1),), "weight", init_ws)?; + Ok(PReLU::new(ws, num_channels.is_none())) }