Avoid generating variables for RNNs without biases. (LaurentMazare#303)

Hakuyume · Dec 18, 2020 · c30de83 · c30de83
1 parent c298235
commit c30de83
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 8 deletions.
diff --git a/src/nn/rnn.rs b/src/nn/rnn.rs
@@ -102,12 +102,14 @@ pub fn lstm(vs: &super::var_store::Path, in_dim: i64, hidden_dim: i64, c: RNNCon
                 &format!("weight_hh_l{}{}", layer_idx, suffix),
                 &[gate_dim, hidden_dim],
             );
-            let b_ih = vs.zeros(&format!("bias_ih_l{}{}", layer_idx, suffix), &[gate_dim]);
-            let b_hh = vs.zeros(&format!("bias_hh_l{}{}", layer_idx, suffix), &[gate_dim]);
             flat_weights.push(w_ih);
             flat_weights.push(w_hh);
-            flat_weights.push(b_ih);
-            flat_weights.push(b_hh);
+            if c.has_biases {
+                let b_ih = vs.zeros(&format!("bias_ih_l{}{}", layer_idx, suffix), &[gate_dim]);
+                let b_hh = vs.zeros(&format!("bias_hh_l{}{}", layer_idx, suffix), &[gate_dim]);
+                flat_weights.push(b_ih);
+                flat_weights.push(b_hh);
+            }
         }
     }
     if vs.device().is_cuda() && crate::Cuda::cudnn_is_available() {
@@ -206,12 +208,14 @@ pub fn gru(vs: &super::var_store::Path, in_dim: i64, hidden_dim: i64, c: RNNConf
                 &format!("weight_hh_l{}{}", layer_idx, suffix),
                 &[gate_dim, hidden_dim],
             );
-            let b_ih = vs.zeros(&format!("bias_ih_l{}{}", layer_idx, suffix), &[gate_dim]);
-            let b_hh = vs.zeros(&format!("bias_hh_l{}{}", layer_idx, suffix), &[gate_dim]);
             flat_weights.push(w_ih);
             flat_weights.push(w_hh);
-            flat_weights.push(b_ih);
-            flat_weights.push(b_hh);
+            if c.has_biases {
+                let b_ih = vs.zeros(&format!("bias_ih_l{}{}", layer_idx, suffix), &[gate_dim]);
+                let b_hh = vs.zeros(&format!("bias_hh_l{}{}", layer_idx, suffix), &[gate_dim]);
+                flat_weights.push(b_ih);
+                flat_weights.push(b_hh);
+            }
         }
     }
     if vs.device().is_cuda() && crate::Cuda::cudnn_is_available() {

diff --git a/tests/nn_tests.rs b/tests/nn_tests.rs
@@ -165,6 +165,10 @@ fn gru_test(rnn_config: nn::RNNConfig) {
 #[test]
 fn gru() {
     gru_test(Default::default());
+    gru_test(nn::RNNConfig {
+        has_biases: false,
+        ..Default::default()
+    });
     gru_test(nn::RNNConfig {
         bidirectional: true,
         ..Default::default()
@@ -210,6 +214,10 @@ fn lstm_test(rnn_config: nn::RNNConfig) {
 #[test]
 fn lstm() {
     lstm_test(Default::default());
+    lstm_test(nn::RNNConfig {
+        has_biases: false,
+        ..Default::default()
+    });
     lstm_test(nn::RNNConfig {
         bidirectional: true,
         ..Default::default()