From 6331bccefc2094ab1f4d2d4056671a88b18b637d Mon Sep 17 00:00:00 2001
From: Eren Golge <egolge@mozilla.com>
Date: Sun, 12 May 2019 17:35:31 +0200
Subject: [PATCH] make dropout oprional #2

---
 layers/tacotron2.py    | 6 +++---
 models/tacotron2.py    | 4 ++--
 utils/generic_utils.py | 1 +
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/layers/tacotron2.py b/layers/tacotron2.py
index 175a8ef3..c0aeafda 100644
--- a/layers/tacotron2.py
+++ b/layers/tacotron2.py
@@ -323,7 +323,7 @@ class Encoder(nn.Module):
 
 # adapted from https://github.com/NVIDIA/tacotron2/
 class Decoder(nn.Module):
-    def __init__(self, in_features, inputs_dim, r, attn_win, attn_norm, prenet_type, forward_attn, trans_agent, location_attn):
+    def __init__(self, in_features, inputs_dim, r, attn_win, attn_norm, prenet_type, prenet_dropout, forward_attn, trans_agent, location_attn):
         super(Decoder, self).__init__()
         self.mel_channels = inputs_dim
         self.r = r
@@ -336,7 +336,7 @@ class Decoder(nn.Module):
         self.p_attention_dropout = 0.1
         self.p_decoder_dropout = 0.1
 
-        self.prenet = Prenet(self.mel_channels * r, prenet_type,
+        self.prenet = Prenet(self.mel_channels * r, prenet_type, prenet_dropout,
                              [self.prenet_dim, self.prenet_dim])
 
         self.attention_rnn = nn.LSTMCell(self.prenet_dim + in_features,
@@ -485,7 +485,7 @@ class Decoder(nn.Module):
             stop_flags[2] = t > inputs.shape[1] * 2
             if all(stop_flags):
                 stop_count += 1
-                if stop_count > 2:
+                if stop_count > 5:
                     break
             elif len(outputs) == self.max_decoder_steps:
                 print("   | > Decoder stopped with 'max_decoder_steps")
diff --git a/models/tacotron2.py b/models/tacotron2.py
index c492c7b1..e9ce1a1b 100644
--- a/models/tacotron2.py
+++ b/models/tacotron2.py
@@ -9,7 +9,7 @@ from utils.generic_utils import sequence_mask
 
 # TODO: match function arguments with tacotron
 class Tacotron2(nn.Module):
-    def __init__(self, num_chars, r, attn_win=False, attn_norm="softmax", prenet_type="original", forward_attn=False, trans_agent=False, location_attn=True):
+    def __init__(self, num_chars, r, attn_win=False, attn_norm="softmax", prenet_type="original", prenet_dropout=True, forward_attn=False, trans_agent=False, location_attn=True):
         super(Tacotron2, self).__init__()
         self.n_mel_channels = 80
         self.n_frames_per_step = r
@@ -18,7 +18,7 @@ class Tacotron2(nn.Module):
         val = sqrt(3.0) * std  # uniform bounds for std
         self.embedding.weight.data.uniform_(-val, val)
         self.encoder = Encoder(512)
-        self.decoder = Decoder(512, self.n_mel_channels, r, attn_win, attn_norm, prenet_type, forward_attn, trans_agent, location_attn)
+        self.decoder = Decoder(512, self.n_mel_channels, r, attn_win, attn_norm, prenet_type, prenet_dropout, forward_attn, trans_agent, location_attn)
         self.postnet = Postnet(self.n_mel_channels)
 
     def shape_outputs(self, mel_outputs, mel_outputs_postnet, alignments):
diff --git a/utils/generic_utils.py b/utils/generic_utils.py
index 19d93888..902affba 100644
--- a/utils/generic_utils.py
+++ b/utils/generic_utils.py
@@ -262,6 +262,7 @@ def setup_model(num_chars, c):
             attn_win=c.windowing,
             attn_norm=c.attention_norm,
             prenet_type=c.prenet_type,
+            prenet_dropout=c.prenet_dropout,
             forward_attn=c.use_forward_attn,
             trans_agent=c.transition_agent,
             location_attn=c.location_attn)