model merge fixes 1
# limitations under the License.
# ==============================================================================
"""
-Modified by blisc to enable support for tacotron models
Custom Helper class that implements the tacotron decoder pre and post nets
"""
from __future__ import absolute_import, division, print_function
@@ -25,6 +24,7 @@ from tensorflow.contrib.seq2seq.python.ops import decoder
from tensorflow.contrib.seq2seq.python.ops.helper import Helper
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
@@ -54,6 +54,8 @@ class TacotronTrainingHelper(Helper):
sequence_length,
prenet=None,
sampling_prob=0.,
+ anneal_teacher_forcing=False,
+ stop_gradient=False,
time_major=False,
sample_ids_shape=None,
sample_ids_dtype=None,
@@ -68,6 +70,8 @@ class TacotronTrainingHelper(Helper):
prenet: prenet to use, currently disabled and used in tacotron decoder
instead.
sampling_prob (float): see tacotron 2 decoder
+ anneal_teacher_forcing (float): see tacotron 2 decoder
+ stop_gradient (float): see tacotron 2 decoder
time_major (bool): (float): see tacotron 2 decoder
mask_decoder_sequence (bool): whether to pass finished when the decoder
passed the sequence_length input or to pass unfinished to dynamic_decode
@@ -82,6 +86,8 @@ class TacotronTrainingHelper(Helper):
self._batch_size = array_ops.size(sequence_length)
self._seed = None
self._sampling_prob = sampling_prob
+ self._anneal_teacher_forcing = anneal_teacher_forcing
+ self._stop_gradient = stop_gradient
self._mask_decoder_sequence = mask_decoder_sequence
self._prenet = prenet
self._zero_inputs = nest.map_structure(
@@ -127,12 +133,13 @@ class TacotronTrainingHelper(Helper):
def get_next_input(inp, out):
next_input = inp.read(time)
+ if self._stop_gradient:
+ next_input = tf.stop_gradient(next_input)
+ out = tf.stop_gradient(out)
if self._prenet is not None:
next_input = self._prenet(next_input)
out = self._prenet(out)
- if self._sampling_prob > 0.:
- next_input = tf.stop_gradient(next_input)
- out = tf.stop_gradient(out)
+ if self._anneal_teacher_forcing or self._sampling_prob > 0:
select_sampler = bernoulli.Bernoulli(
probs=self._sampling_prob, dtype=dtypes.bool
)
