From 4079ec379cb93dff8aadd1ad585252e01a3b44e4 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 30 Dec 2019 03:41:57 -0500 Subject: [PATCH] Tacotron model JIT works? --- tacotron_melgan/tacotron_model.py | 7 ++++--- tacotron_melgan/wrap_tacotron2_model.py | 13 ++++++++----- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/tacotron_melgan/tacotron_model.py b/tacotron_melgan/tacotron_model.py index 5f0419c279..b42e5d74a6 100644 --- a/tacotron_melgan/tacotron_model.py +++ b/tacotron_melgan/tacotron_model.py @@ -255,8 +255,8 @@ class Decoder(nn.Module): decoder_input: all zeros frames """ B = memory.size(0) - decoder_input = Variable(memory.data.new( - B, self.n_mel_channels * self.n_frames_per_step).zero_()) + mem = memory.data.new(B, self.n_mel_channels * self.n_frames_per_step) + decoder_input = Variable(mem.zero_()) return decoder_input def initialize_decoder_states(self, memory, mask): @@ -548,4 +548,5 @@ class Tacotron2(nn.Module): alignments ]) - return outputs + #return outputs + return outputs[0] diff --git a/tacotron_melgan/wrap_tacotron2_model.py b/tacotron_melgan/wrap_tacotron2_model.py index 69c2064464..a9c12da25e 100755 --- a/tacotron_melgan/wrap_tacotron2_model.py +++ b/tacotron_melgan/wrap_tacotron2_model.py @@ -76,15 +76,18 @@ module = Tacotron2(HParams()) print('Load state dict...') module.load_state_dict(tacotron_model['state_dict']) -#module.eval(inference=False) +module.eval() # NB: Complains unless called `Did you forget call .eval() on your model?` -print('JIT model...') -mel_file = '/home/bt/dev/voder/data/mels/LJ002-0320.mel' -example = torch.load(mel_file, map_location=torch.device('cpu')) # TODO -traced_script_module = torch.jit.trace(module, example) +#print('JIT model...') +text_sequence_file = '/home/bt/dev/voder/data/text/tacotron_text_sequence.pt' +#text_sequence_file = '/home/bt/dev/voder/data/mels/LJ002-0320.mel' +text_sequence = torch.load(text_sequence_file, map_location=torch.device('cpu')) +# NB: Getting `Tracing failed sanity checks! Graphs differed across invocations!` +traced_script_module = torch.jit.trace(module, text_sequence, check_trace=False) traced_script_module.save("tacotron_container.pt") #print('Saving model...') #container = torch.jit.script(module) #container.save("tacotron_container.pt") +print('Done.')