huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use the proper broadcasting for prelu. (#1406)

Browse Source
This commit is contained in:
Laurent Mazare
2023-12-05 07:09:31 +01:00
committed by GitHub
parent b5c283e86f
commit 2648e797c2
1 changed files with 16 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
candle-nn/src/activation.rs
View File

@ -65,6 +65,15 @@ impl candle::Module for PReLU {
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
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<usize>, vs: crate::VarBuilder) -> Result<PReLU> {
/// * `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<usize>, vs: crate::VarBuilder) -> Result<PReLU> {
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()))
}