Set attention norm method by config.json

This commit is contained in:
Eren Golge
2019-03-26 00:48:12 +01:00
parent 786510cd6a
commit 0a92c6d5a7
8 changed files with 31 additions and 16 deletions
+3 -2
View File
@@ -14,7 +14,8 @@ class Tacotron(nn.Module):
r=5,
padding_idx=None,
memory_size=5,
attn_windowing=False):
attn_windowing=False,
attn_norm="sigmoid"):
super(Tacotron, self).__init__()
self.r = r
self.mel_dim = mel_dim
@@ -22,7 +23,7 @@ class Tacotron(nn.Module):
self.embedding = nn.Embedding(num_chars, 256, padding_idx=padding_idx)
self.embedding.weight.data.normal_(0, 0.3)
self.encoder = Encoder(256)
self.decoder = Decoder(256, mel_dim, r, memory_size, attn_windowing)
self.decoder = Decoder(256, mel_dim, r, memory_size, attn_windowing, attn_norm)
self.postnet = PostCBHG(mel_dim)
self.last_linear = nn.Sequential(
nn.Linear(self.postnet.cbhg.gru_features * 2, linear_dim),
+2 -2
View File
@@ -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):
def __init__(self, num_chars, r, attn_win=False, attn_norm="softmax"):
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)
self.decoder = Decoder(512, self.n_mel_channels, r, attn_win, attn_norm)
self.postnet = Postnet(self.n_mel_channels)
def shape_outputs(self, mel_outputs, mel_outputs_postnet, alignments):