Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add/lstm direction #2455

Merged
merged 6 commits into from
Sep 30, 2024
Merged

Add/lstm direction #2455

Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d7cf041
add: direction for lstm layer
singjc Aug 29, 2024
afdd4bf
lint: remove unused Error import
singjc Aug 29, 2024
fb68d2a
refactor: remove unnecessary int assignment to Direction enum:
singjc Sep 18, 2024
9ccdde8
refactor: use &'static str type instead of String for direction_str:
singjc Sep 18, 2024
7e1a6d8
Run cargofmt.
LaurentMazare Sep 30, 2024
58e5d6c
Merge remote-tracking branch 'origin/main' into add/lstm_direction
LaurentMazare Sep 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions candle-nn/src/rnn.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,12 @@ impl LSTMState {
}
}

#[derive(Debug, Clone, Copy)]
pub enum Direction {
Forward = 1,
singjc marked this conversation as resolved.
Show resolved Hide resolved
Backward = 2,
}

#[allow(clippy::upper_case_acronyms)]
#[derive(Debug, Clone, Copy)]
pub struct LSTMConfig {
Expand All @@ -78,6 +84,7 @@ pub struct LSTMConfig {
pub b_ih_init: Option<super::Init>,
pub b_hh_init: Option<super::Init>,
pub layer_idx: usize,
pub direction: Direction,
}

impl Default for LSTMConfig {
Expand All @@ -88,6 +95,7 @@ impl Default for LSTMConfig {
b_ih_init: Some(super::Init::Const(0.)),
b_hh_init: Some(super::Init::Const(0.)),
layer_idx: 0,
direction: Direction::Forward,
}
}
}
Expand All @@ -100,6 +108,7 @@ impl LSTMConfig {
b_ih_init: None,
b_hh_init: None,
layer_idx: 0,
direction: Direction::Forward,
}
}
}
Expand Down Expand Up @@ -128,25 +137,29 @@ pub fn lstm(
vb: crate::VarBuilder,
) -> Result<LSTM> {
let layer_idx = config.layer_idx;
let direction_str = match config.direction {
Direction::Forward => String::from(""),
singjc marked this conversation as resolved.
Show resolved Hide resolved
Direction::Backward => String::from("_reverse"),
};
let w_ih = vb.get_with_hints(
(4 * hidden_dim, in_dim),
&format!("weight_ih_l{layer_idx}"), // Only a single layer is supported.
&format!("weight_ih_l{layer_idx}{direction_str}"), // Only a single layer is supported.
config.w_ih_init,
)?;
let w_hh = vb.get_with_hints(
(4 * hidden_dim, hidden_dim),
&format!("weight_hh_l{layer_idx}"), // Only a single layer is supported.
&format!("weight_hh_l{layer_idx}{direction_str}"), // Only a single layer is supported.
config.w_hh_init,
)?;
let b_ih = match config.b_ih_init {
Some(init) => {
Some(vb.get_with_hints(4 * hidden_dim, &format!("bias_ih_l{layer_idx}"), init)?)
Some(vb.get_with_hints(4 * hidden_dim, &format!("bias_ih_l{layer_idx}{direction_str}"), init)?)
}
None => None,
};
let b_hh = match config.b_hh_init {
Some(init) => {
Some(vb.get_with_hints(4 * hidden_dim, &format!("bias_hh_l{layer_idx}"), init)?)
Some(vb.get_with_hints(4 * hidden_dim, &format!("bias_hh_l{layer_idx}{direction_str}"), init)?)
}
None => None,
};
Expand Down