67809715652a92b22870c50ad30f6ff38e292006aedc75ddbdc828aa856ef68f

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,12 +1,71 @@
----
-title: Tts Rvc Autopst
-emoji: 💬
-colorFrom: yellow
-colorTo: purple
-sdk: gradio
-sdk_version: 4.36.1
-app_file: app.py
-pinned: false
----
-
-An example chatbot using [Gradio](https://gradio.app), [`huggingface_hub`](https://huggingface.co/docs/huggingface_hub/v0.22.2/en/index), and the [Hugging Face Inference API](https://huggingface.co/docs/api-inference/index).
+# Global Prosody Style Transfer Without Text Transcriptions
+
+This repository provides a PyTorch implementation of [AutoPST](https://arxiv.org/abs/2106.08519), which enables unsupervised global prosody conversion without text transcriptions.
+
+This is a short video that explains the main concepts of our work. If you find this work useful and use it in your research, please consider citing our paper.
+
+[![SpeechSplit](./assets/cover.png)](https://youtu.be/wow2DRuJ69c/)
+
+```
+@InProceedings{pmlr-v139-qian21b,
+  title = 	 {Global Prosody Style Transfer Without Text Transcriptions},
+  author =       {Qian, Kaizhi and Zhang, Yang and Chang, Shiyu and Xiong, Jinjun and Gan, Chuang and Cox, David and Hasegawa-Johnson, Mark},
+  booktitle = 	 {Proceedings of the 38th International Conference on Machine Learning},
+  pages = 	 {8650--8660},
+  year = 	 {2021},
+  editor = 	 {Meila, Marina and Zhang, Tong},
+  volume = 	 {139},
+  series = 	 {Proceedings of Machine Learning Research},
+  month = 	 {18--24 Jul},
+  publisher =    {PMLR},
+  url = 	 {http://proceedings.mlr.press/v139/qian21b.html}
+}
+
+```
+
+
+## Audio Demo
+
+The audio demo for AutoPST can be found [here](https://auspicious3000.github.io/AutoPST-Demo/)
+
+## Dependencies
+- Python 3.6
+- Numpy
+- Scipy
+- PyTorch == v1.6.0
+- librosa
+- pysptk
+- soundfile
+- wavenet_vocoder ```pip install wavenet_vocoder==0.1.1```
+  for more information, please refer to https://github.com/r9y9/wavenet_vocoder
+
+
+## To Run Demo
+
+Download [pre-trained models](https://drive.google.com/file/d/1ji3Bk6YGvXkPqFu1hLOAJp_SKw-vHGrp/view?usp=sharing) to ```assets```
+
+Download the same WaveNet vocoder model as in [AutoVC](https://github.com/auspicious3000/autovc) to ```assets```
+
+The fast and high-quality hifi-gan v1 (https://github.com/jik876/hifi-gan) pre-trained model is now available [here.](https://drive.google.com/file/d/1n76jHs8k1sDQ3Eh5ajXwdxuY_EZw4N9N/view?usp=sharing)
+
+Please refer to [AutoVC](https://github.com/auspicious3000/autovc) if you have any problems with the vocoder part, because they share the same vocoder scripts.
+
+Run ```demo.ipynb``` 
+
+
+## To Train
+
+Download [training data](https://drive.google.com/file/d/1H1dyA80qREKLHybqnYaqBRRsacIdFbnE/view?usp=sharing) to ```assets```.
+The provided training data is very small for code verification purpose only.
+Please use the scripts to prepare your own data for training.
+
+1. Prepare training data: ```python prepare_train_data.py```
+
+2. Train 1st Stage: ```python main_1.py```
+
+3. Train 2nd Stage: ```python main_2.py```
+
+
+## Final Words
+
+This project is part of an ongoing research. We hope this repo is useful for your research. If you need any help or have any suggestions on improving the framework, please raise an issue and we will do our best to get back to you as soon as possible.

assets/mfcc_stats.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e34f5e00e6eb8591e8dcc3796a56a048c8512245ca552c83069a4c8eb3a57387
+size 52719

assets/spk2emb_82.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea97c16a4d1d2eca10a0481630daa8590d9d8a1e27a87e53fa4340c17440745a
+size 32133

data_loader.py

@@ -0,0 +1,228 @@
+import os
+import pickle 
+import torch
+import numpy as np
+       
+from numpy.random import uniform
+from torch.utils import data
+from torch.utils.data.sampler import Sampler
+from multiprocessing import Process, Manager  
+
+
+
+class Utterances(data.Dataset):
+    """Dataset class for the Utterances dataset."""
+
+    def __init__(self, hparams):
+        """Initialize and preprocess the Utterances dataset."""
+        self.meta_file = hparams.meta_file
+        
+        self.feat_dir_1 = hparams.feat_dir_1
+        self.feat_dir_2 = hparams.feat_dir_2
+        self.feat_dir_3 = hparams.feat_dir_3
+        
+        self.step = 4
+        self.split = 0
+         
+        self.max_len_pad = hparams.max_len_pad
+        
+        meta = pickle.load(open(self.meta_file, "rb"))
+        
+        manager = Manager()
+        meta = manager.list(meta)
+        dataset = manager.list(len(meta)*[None])  # <-- can be shared between processes.
+        processes = []
+        for i in range(0, len(meta), self.step):
+            p = Process(target=self.load_data, 
+                        args=(meta[i:i+self.step],dataset,i))  
+            p.start()
+            processes.append(p)
+        for p in processes:
+            p.join()
+            
+        # very importtant to do dataset = list(dataset)            
+        self.train_dataset = list(dataset)
+        self.num_tokens = len(self.train_dataset)
+        
+        print('Finished loading the {} Utterances training dataset...'.format(self.num_tokens))
+        
+        
+    def load_data(self, submeta, dataset, idx_offset):  
+        for k, sbmt in enumerate(submeta):    
+            uttrs = len(sbmt)*[None]
+            for j, tmp in enumerate(sbmt):
+                if j < 2: 
+                    # fill in speaker name and embedding
+                    uttrs[j] = tmp
+                else:    
+                    # fill in data
+                    sp_tmp = np.load(os.path.join(self.feat_dir_1, tmp))
+                    cep_tmp = np.load(os.path.join(self.feat_dir_2, tmp))[:, 0:14]
+                    cd_tmp = np.load(os.path.join(self.feat_dir_3, tmp))
+                    
+                    assert len(sp_tmp) == len(cep_tmp) == len(cd_tmp)           
+
+                    uttrs[j] = ( np.clip(sp_tmp, 0, 1), cep_tmp, cd_tmp )
+            dataset[idx_offset+k] = uttrs
+    
+    
+    def segment_np(self, cd_long, tau=2):
+        
+        cd_norm = np.sqrt((cd_long ** 2).sum(axis=-1, keepdims=True))
+        G = (cd_long @ cd_long.T) / (cd_norm @ cd_norm.T)
+        
+        L = G.shape[0]
+        
+        num_rep = []
+        num_rep_sync = []
+        
+        prev_boundary = 0
+        rate = np.random.uniform(0.8, 1.3)
+        
+        for t in range(1, L+1):
+            if t==L:
+                num_rep.append(t - prev_boundary)
+                num_rep_sync.append(t - prev_boundary)
+                prev_boundary = t
+            if t < L:
+                q = np.random.uniform(rate-0.1, rate)
+                tmp = G[prev_boundary, max(prev_boundary-20, 0):min(prev_boundary+20, L)]
+                if q <= 1:
+                    epsilon = np.quantile(tmp, q)
+                    if np.all(G[prev_boundary, t:min(t+tau, L)] < epsilon):
+                        num_rep.append(t - prev_boundary)
+                        num_rep_sync.append(t - prev_boundary)
+                        prev_boundary = t
+                else:
+                    epsilon = np.quantile(tmp, 2-q)
+                    if np.all(G[prev_boundary, t:min(t+tau, L)] < epsilon):
+                        num_rep.append(t - prev_boundary)    
+                    else:
+                        num_rep.extend([t-prev_boundary-0.5, 0.5])
+                        
+                    num_rep_sync.append(t - prev_boundary)    
+                    prev_boundary = t
+                    
+        num_rep = np.array(num_rep)
+        num_rep_sync = np.array(num_rep_sync)
+        
+        return num_rep, num_rep_sync
+            
+        
+    def __getitem__(self, index):
+        """Return M uttrs for one spkr."""
+        dataset = self.train_dataset
+        
+        list_uttrs = dataset[index]
+        
+        emb_org = list_uttrs[1]
+        
+        uttr = np.random.randint(2, len(list_uttrs))
+        melsp, melcep, cd_real = list_uttrs[uttr]
+        
+        num_rep, num_rep_sync = self.segment_np(cd_real)
+        
+        return melsp, melcep, cd_real, num_rep, num_rep_sync, len(melsp), len(num_rep), len(num_rep_sync), emb_org
+    
+
+    def __len__(self):
+        """Return the number of spkrs."""
+        return self.num_tokens
+    
+    
+
+class MyCollator(object):
+    def __init__(self, hparams):
+        self.max_len_pad = hparams.max_len_pad
+        
+    def __call__(self, batch):
+        new_batch = []
+        
+        l_short_max = 0
+        l_short_sync_max = 0
+        l_real_max = 0
+        
+        for token in batch:
+            sp_real, cep_real, cd_real, rep, rep_sync, l_real, l_short, l_short_sync, emb = token
+            
+            if l_short > l_short_max:
+                l_short_max = l_short  
+                
+            if l_short_sync > l_short_sync_max:
+                l_short_sync_max = l_short_sync
+            
+            if l_real > l_real_max:
+                l_real_max = l_real
+            
+            sp_real_pad = np.pad(sp_real, ((0,self.max_len_pad-l_real),(0,0)), 'constant')
+            cep_real_pad = np.pad(cep_real, ((0,self.max_len_pad-l_real),(0,0)), 'constant')
+            cd_real_pad = np.pad(cd_real, ((0,self.max_len_pad-l_real),(0,0)), 'constant')
+            
+            rep_pad = np.pad(rep, (0,self.max_len_pad-l_short), 'constant')
+            rep_sync_pad = np.pad(rep_sync, (0,self.max_len_pad-l_short_sync), 'constant')
+            
+            new_batch.append( (sp_real_pad, cep_real_pad, cd_real_pad, rep_pad, rep_sync_pad, l_real, l_short, l_short_sync, emb) ) 
+            
+        batch = new_batch  
+        
+        a, b, c, d, e, f, g, h, i = zip(*batch)
+        
+        sp_real = torch.from_numpy(np.stack(a, axis=0))[:,:l_real_max+1,:]
+        cep_real = torch.from_numpy(np.stack(b, axis=0))[:,:l_real_max+1,:]
+        cd_real = torch.from_numpy(np.stack(c, axis=0))[:,:l_real_max+1,:]
+        num_rep = torch.from_numpy(np.stack(d, axis=0))[:,:l_short_max+1]
+        num_rep_sync = torch.from_numpy(np.stack(e, axis=0))[:,:l_short_sync_max+1]
+        
+        len_real = torch.from_numpy(np.stack(f, axis=0))
+        len_short = torch.from_numpy(np.stack(g, axis=0))
+        len_short_sync = torch.from_numpy(np.stack(h, axis=0))
+        
+        spk_emb = torch.from_numpy(np.stack(i, axis=0))
+        
+        return sp_real, cep_real, cd_real, num_rep, num_rep_sync, len_real, len_short, len_short_sync, spk_emb
+
+
+    
+class MultiSampler(Sampler):
+    """Samples elements more than once in a single pass through the data.
+    """
+    def __init__(self, num_samples, n_repeats, shuffle=False):
+        self.num_samples = num_samples
+        self.n_repeats = n_repeats
+        self.shuffle = shuffle
+
+    def gen_sample_array(self):
+        self.sample_idx_array = torch.arange(self.num_samples, dtype=torch.int64).repeat(self.n_repeats)
+        if self.shuffle:
+            self.sample_idx_array = self.sample_idx_array[torch.randperm(len(self.sample_idx_array))]
+        return self.sample_idx_array
+
+    def __iter__(self):
+        return iter(self.gen_sample_array())
+
+    def __len__(self):
+        return len(self.sample_idx_array)        
+    
+    
+    
+def worker_init_fn(x):
+    return np.random.seed((torch.initial_seed()) % (2**32))    
+
+def get_loader(hparams):
+    """Build and return a data loader."""
+    
+    dataset = Utterances(hparams)
+    
+    my_collator = MyCollator(hparams)
+    
+    sampler = MultiSampler(len(dataset), hparams.samplier, shuffle=hparams.shuffle)
+    
+    data_loader = data.DataLoader(dataset=dataset,
+                                  batch_size=hparams.batch_size,
+                                  sampler=sampler,
+                                  num_workers=hparams.num_workers,
+                                  drop_last=True,
+                                  pin_memory=False,
+                                  worker_init_fn=worker_init_fn,
+                                  collate_fn=my_collator)
+    return data_loader

demo.ipynb

@@ -0,0 +1,115 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pickle\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "from collections import OrderedDict\n",
+    "from onmt_modules.misc import sequence_mask\n",
+    "from model_autopst import Generator_2 as Predictor\n",
+    "from hparams_autopst import hparams\n",
+    "\n",
+    "device = 'cuda:0'\n",
+    "\n",
+    "P = Predictor(hparams).eval().to(device)\n",
+    "\n",
+    "checkpoint = torch.load('./assets/580000-P.ckpt', map_location=lambda storage, loc: storage)  \n",
+    "P.load_state_dict(checkpoint['model'], strict=True)\n",
+    "print('Loaded predictor .....................................................')\n",
+    "\n",
+    "dict_test = pickle.load(open('./assets/test_vctk.meta', 'rb'))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "spect_vc = OrderedDict()\n",
+    "\n",
+    "uttrs = [('p231', 'p270', '001'),\n",
+    "         ('p270', 'p231', '001'),\n",
+    "         ('p231', 'p245', '003001'),\n",
+    "         ('p245', 'p231', '003001'),\n",
+    "         ('p239', 'p270', '024002'),\n",
+    "         ('p270', 'p239', '024002')]\n",
+    "\n",
+    "\n",
+    "for uttr in uttrs:\n",
+    "        \n",
+    "    cep_real, spk_emb = dict_test[uttr[0]][uttr[2]]\n",
+    "    cep_real_A = torch.from_numpy(cep_real).unsqueeze(0).to(device)\n",
+    "    len_real_A = torch.tensor(cep_real_A.size(1)).unsqueeze(0).to(device)\n",
+    "    real_mask_A = sequence_mask(len_real_A, cep_real_A.size(1)).float()\n",
+    "    \n",
+    "    _, spk_emb = dict_test[uttr[1]][uttr[2]]\n",
+    "    spk_emb_B = torch.from_numpy(spk_emb).unsqueeze(0).to(device)\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        spect_output, len_spect = P.infer_onmt(cep_real_A.transpose(2,1)[:,:14,:],\n",
+    "                                               real_mask_A,\n",
+    "                                               len_real_A,\n",
+    "                                               spk_emb_B)\n",
+    "    \n",
+    "    uttr_tgt = spect_output[:len_spect[0],0,:].cpu().numpy()\n",
+    "        \n",
+    "    spect_vc[f'{uttr[0]}_{uttr[1]}_{uttr[2]}'] = uttr_tgt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# spectrogram to waveform\n",
+    "# Feel free to use other vocoders\n",
+    "# This cell requires some preparation to work, please see the corresponding part in AutoVC\n",
+    "import torch\n",
+    "import librosa\n",
+    "import pickle\n",
+    "import os\n",
+    "from synthesis import build_model\n",
+    "from synthesis import wavegen\n",
+    "\n",
+    "model = build_model().to(device)\n",
+    "checkpoint = torch.load(\"./assets/checkpoint_step001000000_ema.pth\")\n",
+    "model.load_state_dict(checkpoint[\"state_dict\"])\n",
+    "\n",
+    "for name, sp in spect_vc.items():\n",
+    "    print(name)\n",
+    "    waveform = wavegen(model, c=sp)   \n",
+    "    librosa.output.write_wav('./assets/'+name+'.wav', waveform, sr=16000)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

fast_decoders.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from onmt_modules.misc import sequence_mask
+
+
+class DecodeFunc_Sp(object):
+    """
+    Decoding functions
+    """
+    def __init__(self, hparams, type_out):
+                
+        if type_out == 'Sp':
+            self.dim_freq = hparams.dim_freq
+            self.max_decoder_steps = hparams.dec_steps_sp
+        elif type_out == 'Tx':
+            self.dim_freq = hparams.dim_code
+            self.max_decoder_steps = hparams.dec_steps_tx
+        else:
+            raise ValueError
+        
+        self.gate_threshold = hparams.gate_threshold
+        self.type_out = type_out
+        
+    def __call__(self, tgt, memory_bank, memory_lengths, decoder, postnet):
+        
+        dec_outs, attns = decoder(tgt, memory_bank, step=None, 
+                                  memory_lengths=memory_lengths)
+        spect_gate = postnet(dec_outs)
+        spect, gate = spect_gate[:, :, 1:], spect_gate[:, :, :1]
+        
+        return spect, gate
+    
+    
+    def infer(self, tgt_real, memory_bank, memory_lengths, decoder, postnet):
+        B = memory_bank.size(1)
+        device = memory_bank.device
+        
+        spect_outputs = torch.zeros((self.max_decoder_steps, B, self.dim_freq), 
+                                    dtype=torch.float, device=device)
+        gate_outputs = torch.zeros((self.max_decoder_steps, B, 1), 
+                                   dtype=torch.float, device=device)
+        tgt_words = torch.zeros([B, 1], 
+                                 dtype=torch.float, device=device)
+        
+        current_pred = torch.zeros([1, B, self.dim_freq], 
+                                    dtype=torch.float, device=device)
+        
+        for t in range(self.max_decoder_steps):
+            
+            dec_outs, _ = decoder(current_pred, 
+                                  memory_bank, t, 
+                                  memory_lengths=memory_lengths,
+                                  tgt_words=tgt_words)
+            spect_gate = postnet(dec_outs)
+            
+            spect, gate = spect_gate[:, :, 1:], spect_gate[:, :, :1]
+            spect_outputs[t:t+1] = spect
+            gate_outputs[t:t+1] = gate
+            
+            stop = (torch.sigmoid(gate) - self.gate_threshold + 0.5).round()
+            current_pred = spect.data
+            tgt_words = stop.squeeze(-1).t()
+            
+            if t == self.max_decoder_steps - 1:
+                print(f"Warning! {self.type_out} reached max decoder steps")
+            
+            if (stop == 1).all():
+                break
+        
+        stop_quant = (torch.sigmoid(gate_outputs.data) - self.gate_threshold + 0.5).round().squeeze(-1) 
+        len_spect = (stop_quant.cumsum(dim=0)==0).sum(dim=0)
+        
+        return spect_outputs, len_spect, gate_outputs

hparams_autopst.py

@@ -0,0 +1,62 @@
+from tfcompat.hparam import HParams
+
+# NOTE: If you want full control for model architecture. please take a look
+# at the code and change whatever you want. Some hyper parameters are hardcoded.
+
+# Default hyperparameters:
+hparams = HParams(
+    
+    # sea params
+    dim_neck_sea = 4,
+    dim_freq_sea = 14,
+    dim_enc_sea = 512,
+    
+    # autopst params
+    dim_freq = 80,
+    dim_code = 4,
+    dim_spk = 82,
+    dim_sty = 128,
+    gate_threshold = 0.48,
+    dec_steps_tx = 640,
+    dec_steps_sp = 806,
+    chs_grp = 16,
+    
+    # onmt params
+    enc_layers = 4,
+    enc_rnn_size = 256,
+    dec_layers = 4, 
+    dec_rnn_size = 256, 
+    transformer_ff = 1024,
+    heads = 8, 
+    dropout = 0.1,
+    attention_dropout = 0.1,  
+    max_relative_positions = 0,
+    copy_attn = False, 
+    self_attn_type = "scaled-dot", 
+    aan_useffn = False,
+    full_context_alignment = False, 
+    alignment_layer = 0,
+    alignment_heads = 0,
+    
+    # pretrained model
+    pretrained_path = './assets/xxx.ckpt',
+    
+    # data loader
+    meta_file = './assets/train_vctk.meta',
+    feat_dir_1 = './assets/vctk16-train-sp-mel',
+    feat_dir_2 = './assets/vctk16-train-cep-mel',
+    feat_dir_3 = './assets/vctk16-train-teacher', 
+    batch_size = 4,
+    shuffle = True,
+    num_workers = 0,
+    samplier = 2,
+    max_len_pad = 2048,
+    sampling_params = (0.8, 1.3, 0.1),
+    
+)
+
+
+def hparams_debug_string():
+    values = hparams.values()
+    hp = ['  %s: %s' % (name, values[name]) for name in values]
+    return 'Hyperparameters:\n' + '\n'.join(hp)

hparams_sea.py

@@ -0,0 +1,21 @@
+from tfcompat.hparam import HParams
+
+# NOTE: If you want full control for model architecture. please take a look
+# at the code and change whatever you want. Some hyper parameters are hardcoded.
+
+# Default hyperparameters:
+hparams = HParams(
+    dim_neck_sea = 8,
+    dim_freq_sea = 20,
+    dim_spk = 82,
+    dim_enc_sea = 512,
+    chs_grp = 16,
+    dim_freq_sp = 80,
+  
+)
+
+
+def hparams_debug_string():
+    values = hparams.values()
+    hp = ['  %s: %s' % (name, values[name]) for name in values]
+    return 'Hyperparameters:\n' + '\n'.join(hp)

main_1.py

@@ -0,0 +1,41 @@
+import os
+import argparse
+import torch
+
+from solver_1 import Solver
+from data_loader import get_loader
+from hparams_autopst import hparams, hparams_debug_string
+
+
+
+def str2bool(v):
+    return v.lower() in ('true')
+
+def main(config):
+
+    # Data loader
+    data_loader = get_loader(hparams)
+    
+    # Solver for training 
+    solver = Solver(data_loader, config, hparams)
+
+    solver.train()
+    
+        
+    
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    
+    # Training configuration.
+    parser.add_argument('--num_iters', type=int, default=1000000)
+
+    # Miscellaneous.
+    parser.add_argument('--device_id', type=int, default=0)
+
+    # Step size.
+    parser.add_argument('--log_step', type=int, default=10)
+    
+    config = parser.parse_args()
+    print(config)
+    print(hparams_debug_string())
+    main(config)

main_2.py

@@ -0,0 +1,41 @@
+import os
+import argparse
+import torch
+
+from solver_2 import Solver
+from data_loader import get_loader
+from hparams_autopst import hparams, hparams_debug_string
+
+
+
+def str2bool(v):
+    return v.lower() in ('true')
+
+def main(config):
+
+    # Data loader
+    data_loader = get_loader(hparams)
+    
+    # Solver for training 
+    solver = Solver(data_loader, config, hparams)
+
+    solver.train()
+    
+        
+    
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    
+    # Training configuration.
+    parser.add_argument('--num_iters', type=int, default=1000000)
+
+    # Miscellaneous.
+    parser.add_argument('--device_id', type=int, default=0)
+
+    # Step size.
+    parser.add_argument('--log_step', type=int, default=10)
+    
+    config = parser.parse_args()
+    print(config)
+    print(hparams_debug_string())
+    main(config)

model_autopst.py

@@ -0,0 +1,258 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from utils import filter_bank_mean
+
+from fast_decoders import DecodeFunc_Sp
+
+from model_sea import Encoder_2 as Encoder_Code_2
+
+from override_decoder import OnmtDecoder_1 as OnmtDecoder
+
+from onmt_modules.misc import sequence_mask
+from onmt_modules.embeddings import PositionalEncoding
+from onmt_modules.encoder_transformer import TransformerEncoder as OnmtEncoder
+
+      
+        
+class Prenet(nn.Module):
+    def __init__(self, dim_input, dim_output, dropout=0.1):
+        super().__init__() 
+        
+        mlp = nn.Linear(dim_input, dim_output, bias=True)
+        pe = PositionalEncoding(dropout, dim_output, 1600)
+        
+        self.make_prenet = nn.Sequential()
+        self.make_prenet.add_module('mlp', mlp)
+        self.make_prenet.add_module('pe', pe)
+        
+        self.word_padding_idx = 1
+        
+    def forward(self, source, step=None):
+        
+        for i, module in enumerate(self.make_prenet._modules.values()):
+            if i == len(self.make_prenet._modules.values()) - 1:
+                source = module(source, step=step)
+            else:
+                source = module(source)
+                
+        return source
+    
+    
+    
+class Decoder_Sp(nn.Module):
+    """
+    Speech Decoder
+    """
+    def __init__(self, hparams):
+        super().__init__() 
+        
+        self.dim_freq = hparams.dim_freq
+        self.max_decoder_steps = hparams.dec_steps_sp
+        self.gate_threshold = hparams.gate_threshold
+        
+        prenet = Prenet(hparams.dim_freq, hparams.dec_rnn_size)
+        self.decoder = OnmtDecoder.from_opt(hparams, prenet)
+
+        self.postnet = nn.Linear(hparams.dec_rnn_size, 
+                                 hparams.dim_freq+1, bias=True)
+        
+    def forward(self, tgt, tgt_lengths, memory_bank, memory_lengths):
+        
+        dec_outs, attns = self.decoder(tgt, memory_bank, step=None, 
+                                       memory_lengths=memory_lengths,
+                                       tgt_lengths=tgt_lengths)
+        spect_gate = self.postnet(dec_outs)
+        spect, gate = spect_gate[:, :, 1:], spect_gate[:, :, :1]
+        
+        return spect, gate
+    
+    
+    
+class Encoder_Tx_Spk(nn.Module):
+    """
+    Text Encoder
+    """
+    def __init__(self, hparams):
+        super().__init__() 
+        
+        prenet = Prenet(hparams.dim_code+hparams.dim_spk, 
+                        hparams.enc_rnn_size)
+        self.encoder = OnmtEncoder.from_opt(hparams, prenet)
+        
+    def forward(self, src, src_lengths, spk_emb):
+        
+        spk_emb = spk_emb.unsqueeze(0).expand(src.size(0),-1,-1)
+        src_spk = torch.cat((src, spk_emb), dim=-1)
+        enc_states, memory_bank, src_lengths = self.encoder(src_spk, src_lengths)
+        
+        return enc_states, memory_bank, src_lengths
+    
+    
+    
+class Decoder_Tx(nn.Module):
+    """
+    Text Decoder with stop 
+    and num_rep prediction
+    """
+    def __init__(self, hparams):
+        super().__init__()        
+        
+        self.dim_code = hparams.dim_code
+        self.max_decoder_steps = hparams.dec_steps_tx
+        self.gate_threshold = hparams.gate_threshold
+        self.dim_rep = hparams.dim_rep
+        
+        prenet = Prenet(hparams.dim_code, hparams.dec_rnn_size)
+        self.decoder = OnmtDecoder.from_opt(hparams, prenet)
+
+        self.postnet_1 = nn.Linear(hparams.dec_rnn_size, 
+                                   hparams.dim_code+1, bias=True)
+        
+        self.postnet_2 = nn.Linear(hparams.dec_rnn_size, 
+                                   self.dim_rep, bias=True)
+        
+    def forward(self, tgt, tgt_lengths, memory_bank, memory_lengths):
+        
+        dec_outs, attns = self.decoder(tgt, memory_bank, step=None, 
+                                       memory_lengths=memory_lengths,
+                                       tgt_lengths=tgt_lengths)
+        gate_text = self.postnet_1(dec_outs)
+        rep = self.postnet_2(dec_outs)
+        gate, text = gate_text[:, :, :1], gate_text[:, :, 1:]
+        
+        return text, gate, rep
+        
+        
+    
+class Generator_1(nn.Module):
+    '''
+    sync stage 1
+    '''
+    def __init__(self, hparams):
+        super().__init__() 
+        
+        self.encoder_cd = Encoder_Code_2(hparams)
+        self.encoder_tx = Encoder_Tx_Spk(hparams)
+        self.decoder_sp = Decoder_Sp(hparams)   
+        self.encoder_spk = nn.Linear(hparams.dim_spk, 
+                                     hparams.enc_rnn_size, bias=True)
+        self.fast_dec_sp = DecodeFunc_Sp(hparams, 'Sp')
+        
+        
+    def pad_sequences_rnn(self, cd_short, num_rep, len_long):
+        B, L, C = cd_short.size()
+        out_tensor = torch.zeros((B, len_long.max(), C), device=cd_short.device)
+        '''
+        len_long = len_spect + 1
+        '''
+        for i in range(B):
+            code_sync = cd_short[i].repeat_interleave(num_rep[i], dim=0)
+            out_tensor[i, :len_long[i]-1, :] = code_sync
+            
+        return out_tensor 
+
+        
+    def forward(self, cep_in, mask_long, codes_mask, num_rep, len_short,
+                      tgt_spect, len_spect, 
+                      spk_emb):
+        
+        cd_long = self.encoder_cd(cep_in, mask_long)
+        fb = filter_bank_mean(num_rep, codes_mask, cd_long.size(1))
+        
+        cd_short = torch.bmm(fb.detach(), cd_long)
+        
+        cd_short_sync = self.pad_sequences_rnn(cd_short, num_rep, len_spect)
+        
+        spk_emb_1 = self.encoder_spk(spk_emb)
+        
+        # text to speech
+        _, memory_tx, _ = self.encoder_tx(cd_short_sync.transpose(1,0), len_spect, 
+                                          spk_emb)
+        memory_tx_spk = torch.cat((spk_emb_1.unsqueeze(0), memory_tx), dim=0)
+        self.decoder_sp.decoder.init_state(memory_tx_spk, None, None)
+        spect_out, gate_sp_out \
+        = self.decoder_sp(tgt_spect, len_spect, memory_tx_spk, len_spect+1)
+        
+        return spect_out, gate_sp_out
+    
+    
+    def infer_onmt(self, cep_in, mask_long,
+                   len_spect, 
+                   spk_emb):
+        
+        cd_long = self.encoder_cd(cep_in, mask_long)
+        
+        spk_emb_1 = self.encoder_spk(spk_emb)
+        
+        # text to speech
+        _, memory_tx, _ = self.encoder_tx(cd_long.transpose(1,0), len_spect,
+                                          spk_emb)
+        memory_tx_spk = torch.cat((spk_emb_1.unsqueeze(0), memory_tx), dim=0)
+        self.decoder_sp.decoder.init_state(memory_tx_spk, None, None)
+        spect_output, len_spect_out, stop_sp_output \
+        = self.fast_dec_sp.infer(None, memory_tx_spk, len_spect+1, 
+                                 self.decoder_sp.decoder, 
+                                 self.decoder_sp.postnet)
+        
+        return spect_output, len_spect_out
+    
+    
+    
+class Generator_2(nn.Module):
+    '''
+    async stage 2
+    '''
+    def __init__(self, hparams):
+        super().__init__() 
+        
+        self.encoder_cd = Encoder_Code_2(hparams)
+        self.encoder_tx = Encoder_Tx_Spk(hparams)
+        self.decoder_sp = Decoder_Sp(hparams)   
+        self.encoder_spk = nn.Linear(hparams.dim_spk, 
+                                     hparams.enc_rnn_size, bias=True)
+        self.fast_dec_sp = DecodeFunc_Sp(hparams, 'Sp')
+        
+        
+    def forward(self, cep_in, mask_long, codes_mask, num_rep, len_short,
+                      tgt_spect, len_spect, 
+                      spk_emb):
+        
+        cd_long = self.encoder_cd(cep_in, mask_long)
+        fb = filter_bank_mean(num_rep, codes_mask, cd_long.size(1))
+        
+        cd_short = torch.bmm(fb.detach(), cd_long.detach())
+        
+        spk_emb_1 = self.encoder_spk(spk_emb)
+        
+        # text to speech
+        _, memory_tx, _ = self.encoder_tx(cd_short.transpose(1,0), len_short, 
+                                          spk_emb)
+        memory_tx_spk = torch.cat((spk_emb_1.unsqueeze(0), memory_tx), dim=0)
+        self.decoder_sp.decoder.init_state(memory_tx_spk, None, None)
+        spect_out, gate_sp_out \
+        = self.decoder_sp(tgt_spect, len_spect, memory_tx_spk, len_short+1)
+        
+        return spect_out, gate_sp_out
+    
+    
+    def infer_onmt(self, cep_in, mask_long, len_spect,
+                   spk_emb):
+        
+        cd_long = self.encoder_cd(cep_in, mask_long)
+        
+        spk_emb_1 = self.encoder_spk(spk_emb)
+        
+        # text to speech
+        _, memory_tx, _ = self.encoder_tx(cd_long.transpose(1,0), len_spect, 
+                                          spk_emb)
+        memory_tx_spk = torch.cat((spk_emb_1.unsqueeze(0), memory_tx), dim=0)
+        self.decoder_sp.decoder.init_state(memory_tx_spk, None, None)
+        spect_output, len_spect_out, stop_sp_output \
+        = self.fast_dec_sp.infer(None, memory_tx_spk, len_spect+1, 
+                                 self.decoder_sp.decoder, 
+                                 self.decoder_sp.postnet)
+        
+        return spect_output, len_spect_out

model_sea.py

@@ -0,0 +1,242 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from utils import ConvNorm, LinearNorm
+from torch.nn.parameter import Parameter
+
+
+
+class GroupNorm_Mask(nn.Module):
+    def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
+        super().__init__()
+        
+        self.num_groups = num_groups
+        self.num_channels = num_channels
+        self.eps = eps
+        self.affine = affine
+        if self.affine:
+            self.weight = Parameter(torch.Tensor(num_channels))
+            self.bias = Parameter(torch.Tensor(num_channels))
+        else:
+            self.register_parameter('weight', None)
+            self.register_parameter('bias', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.affine:
+            nn.init.ones_(self.weight)
+            nn.init.zeros_(self.bias)
+    
+    def forward(self, x, mask):
+        B, C, L = x.size()
+        assert C % self.num_groups == 0
+        
+        x = x.view(B, self.num_groups, C//self.num_groups, L)
+        mask = mask.view(B, 1, 1, L)
+        x = x * mask
+        
+        mean = x.mean(dim=2, keepdim=True).sum(dim=3, keepdim=True) / mask.sum(dim=3, keepdim=True)
+        var = (((x - mean)**2)*mask).mean(dim=2, keepdim=True).sum(dim=3, keepdim=True) / mask.sum(dim=3, keepdim=True)
+        
+        x = (x - mean) / (var + self.eps).sqrt()
+        
+        x = x.view(B, C, L)
+        
+        return x * self.weight.view(1,-1,1) + self.bias.view(1,-1,1)
+    
+    
+    
+class M43_Sequential(nn.Sequential):
+    def forward(self, inputs, mask):
+        inputs = self._modules['0'](inputs)
+        inputs = self._modules['1'](inputs, mask)
+        return inputs
+    
+
+
+class Encoder(nn.Module):
+    """Encoder module:
+    """
+    def __init__(self, hparams):
+        super(Encoder, self).__init__()
+        
+        self.dim_freq = hparams.dim_freq_sea
+        self.dim_enc = hparams.dim_enc_sea
+        self.chs_grp = hparams.chs_grp
+        self.dim_neck = hparams.dim_neck_sea
+        
+        convolutions = []        
+        for i in range(5):
+            conv_layer = M43_Sequential(
+                ConvNorm(self.dim_freq if i==0 else self.dim_enc,
+                         self.dim_enc,
+                         kernel_size=1, stride=1,
+                         padding=0,
+                         dilation=1, w_init_gain='relu'),
+                GroupNorm_Mask(self.dim_enc//self.chs_grp, self.dim_enc))
+            convolutions.append(conv_layer)
+             
+        conv_layer = M43_Sequential(
+                ConvNorm(self.dim_enc,
+                         128,
+                         kernel_size=1, stride=1,
+                         padding=0,
+                         dilation=1, w_init_gain='relu'),
+                GroupNorm_Mask(128//self.chs_grp, 128))
+        convolutions.append(conv_layer)   
+        
+        conv_layer = M43_Sequential(
+                ConvNorm(128,
+                         32,
+                         kernel_size=1, stride=1,
+                         padding=0,
+                         dilation=1, w_init_gain='relu'),
+                GroupNorm_Mask(32//self.chs_grp, 32))
+        convolutions.append(conv_layer)           
+        
+        conv_layer = M43_Sequential(
+                ConvNorm(32,
+                         self.dim_neck,
+                         kernel_size=1, stride=1,
+                         padding=0,
+                         dilation=1, w_init_gain='relu'),
+                GroupNorm_Mask(1, self.dim_neck))
+        convolutions.append(conv_layer)   
+            
+        self.convolutions = nn.ModuleList(convolutions)
+        
+
+    def forward(self, x, mask):
+                
+        for conv in self.convolutions:
+            x = F.relu(conv(x, mask))
+            
+        codes = x.permute(0, 2, 1) * mask.unsqueeze(-1)
+
+        return codes
+      
+      
+        
+class Decoder(nn.Module):
+    """Decoder module:
+    """
+    def __init__(self, hparams):
+        super(Decoder, self).__init__()
+        self.dim_enc = hparams.dim_enc_sea
+        self.dim_emb = hparams.dim_spk
+        self.dim_freq = hparams.dim_freq_sp
+        self.dim_neck = hparams.dim_neck_sea
+        
+        self.lstm = nn.LSTM(self.dim_neck+self.dim_emb, 
+                            1024, 3, batch_first=True)
+        
+        self.linear_projection = LinearNorm(1024, self.dim_freq)
+
+    def forward(self, x):
+        
+        outputs = self.lstm(x)[0]
+        
+        decoder_output = self.linear_projection(outputs)
+
+        return decoder_output   
+    
+
+    
+    
+class Generator(nn.Module):
+    """Generator network."""
+    def __init__(self, hparams):
+        super(Generator, self).__init__()
+        
+        self.encoder = Encoder(hparams)
+        self.decoder = Decoder(hparams)
+
+    def forward(self, x, c_trg):
+        
+        x = x.transpose(2,1)
+        codes = self.encoder(x)
+        
+        encoder_outputs = torch.cat((codes, 
+                                     c_trg.unsqueeze(1).expand(-1,x.size(-1),-1)), dim=-1)
+        mel_outputs = self.decoder(encoder_outputs)
+        
+        return mel_outputs
+    
+    def encode(self, x, mask):
+        x = x.transpose(2,1)
+        codes = self.encoder(x, mask)
+        return codes
+    
+    def decode(self, codes, c_trg):
+        encoder_outputs = torch.cat((codes, 
+                                     c_trg.unsqueeze(1).expand(-1,codes.size(1),-1)), dim=-1)
+        mel_outputs = self.decoder(encoder_outputs)
+        
+        return mel_outputs
+    
+    
+    
+class Encoder_2(nn.Module):
+    """Encoder module:
+    """
+    def __init__(self, hparams):
+        super().__init__()
+        
+        self.dim_freq = hparams.dim_freq_sea
+        self.dim_enc = hparams.dim_enc_sea
+        self.chs_grp = hparams.chs_grp
+        self.dim_neck = hparams.dim_neck_sea
+        
+        convolutions = []        
+        for i in range(5):
+            conv_layer = M43_Sequential(
+                ConvNorm(self.dim_freq if i==0 else self.dim_enc,
+                         self.dim_enc,
+                         kernel_size=5, stride=1,
+                         padding=2,
+                         dilation=1, w_init_gain='relu'),
+                GroupNorm_Mask(self.dim_enc//self.chs_grp, self.dim_enc))
+            convolutions.append(conv_layer)
+             
+        conv_layer = M43_Sequential(
+                ConvNorm(self.dim_enc,
+                         128,
+                         kernel_size=5, stride=1,
+                         padding=2,
+                         dilation=1, w_init_gain='relu'),
+                GroupNorm_Mask(128//self.chs_grp, 128))
+        convolutions.append(conv_layer)   
+        
+        conv_layer = M43_Sequential(
+                ConvNorm(128,
+                         32,
+                         kernel_size=5, stride=1,
+                         padding=2,
+                         dilation=1, w_init_gain='relu'),
+                GroupNorm_Mask(32//self.chs_grp, 32))
+        convolutions.append(conv_layer)           
+        
+        conv_layer = M43_Sequential(
+                ConvNorm(32,
+                         self.dim_neck,
+                         kernel_size=5, stride=1,
+                         padding=2,
+                         dilation=1, w_init_gain='linear'),
+                GroupNorm_Mask(1, self.dim_neck))
+        convolutions.append(conv_layer)   
+            
+        self.convolutions = nn.ModuleList(convolutions)
+        
+
+    def forward(self, x, mask):
+                
+        for i in range(len(self.convolutions)-1):
+            x = F.relu(self.convolutions[i](x, mask))
+            
+        x = self.convolutions[-1](x, mask)    
+            
+        codes = x.permute(0, 2, 1) * mask.unsqueeze(-1)
+
+        return codes    

onmt_modules/average_attn.py

@@ -0,0 +1,111 @@
+# -*- coding: utf-8 -*-
+"""Average Attention module."""
+
+import torch
+import torch.nn as nn
+
+from .position_ffn import PositionwiseFeedForward
+
+
+class AverageAttention(nn.Module):
+    """
+    Average Attention module from
+    "Accelerating Neural Transformer via an Average Attention Network"
+    :cite:`DBLP:journals/corr/abs-1805-00631`.
+
+    Args:
+       model_dim (int): the dimension of keys/values/queries,
+           must be divisible by head_count
+       dropout (float): dropout parameter
+    """
+
+    def __init__(self, model_dim, dropout=0.1, aan_useffn=False):
+        self.model_dim = model_dim
+        self.aan_useffn = aan_useffn
+        super(AverageAttention, self).__init__()
+        if aan_useffn:
+            self.average_layer = PositionwiseFeedForward(model_dim, model_dim,
+                                                         dropout)
+        self.gating_layer = nn.Linear(model_dim * 2, model_dim * 2)
+
+    def cumulative_average_mask(self, batch_size, inputs_len, device):
+        """
+        Builds the mask to compute the cumulative average as described in
+        :cite:`DBLP:journals/corr/abs-1805-00631` -- Figure 3
+
+        Args:
+            batch_size (int): batch size
+            inputs_len (int): length of the inputs
+
+        Returns:
+            (FloatTensor):
+
+            * A Tensor of shape ``(batch_size, input_len, input_len)``
+        """
+
+        triangle = torch.tril(torch.ones(inputs_len, inputs_len,
+                              dtype=torch.float, device=device))
+        weights = torch.ones(1, inputs_len, dtype=torch.float, device=device) \
+            / torch.arange(1, inputs_len + 1, dtype=torch.float, device=device)
+        mask = triangle * weights.transpose(0, 1)
+
+        return mask.unsqueeze(0).expand(batch_size, inputs_len, inputs_len)
+
+    def cumulative_average(self, inputs, mask_or_step,
+                           layer_cache=None, step=None):
+        """
+        Computes the cumulative average as described in
+        :cite:`DBLP:journals/corr/abs-1805-00631` -- Equations (1) (5) (6)
+
+        Args:
+            inputs (FloatTensor): sequence to average
+                ``(batch_size, input_len, dimension)``
+            mask_or_step: if cache is set, this is assumed
+                to be the current step of the
+                dynamic decoding. Otherwise, it is the mask matrix
+                used to compute the cumulative average.
+            layer_cache: a dictionary containing the cumulative average
+                of the previous step.
+
+        Returns:
+            a tensor of the same shape and type as ``inputs``.
+        """
+
+        if layer_cache is not None:
+            step = mask_or_step
+            average_attention = (inputs + step *
+                                 layer_cache["prev_g"]) / (step + 1)
+            layer_cache["prev_g"] = average_attention
+            return average_attention
+        else:
+            mask = mask_or_step
+            return torch.matmul(mask.to(inputs.dtype), inputs)
+
+    def forward(self, inputs, mask=None, layer_cache=None, step=None):
+        """
+        Args:
+            inputs (FloatTensor): ``(batch_size, input_len, model_dim)``
+
+        Returns:
+            (FloatTensor, FloatTensor):
+
+            * gating_outputs ``(batch_size, input_len, model_dim)``
+            * average_outputs average attention
+                ``(batch_size, input_len, model_dim)``
+        """
+
+        batch_size = inputs.size(0)
+        inputs_len = inputs.size(1)
+        average_outputs = self.cumulative_average(
+          inputs, self.cumulative_average_mask(batch_size,
+                                               inputs_len, inputs.device)
+          if layer_cache is None else step, layer_cache=layer_cache)
+        if self.aan_useffn:
+            average_outputs = self.average_layer(average_outputs)
+        gating_outputs = self.gating_layer(torch.cat((inputs,
+                                                      average_outputs), -1))
+        input_gate, forget_gate = torch.chunk(gating_outputs, 2, dim=2)
+        gating_outputs = torch.sigmoid(input_gate) * inputs + \
+            torch.sigmoid(forget_gate) * average_outputs
+
+        return gating_outputs, average_outputs

onmt_modules/decoder.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+
+from .misc import aeq
+
+
+class DecoderBase(nn.Module):
+    """Abstract class for decoders.
+
+    Args:
+        attentional (bool): The decoder returns non-empty attention.
+    """
+
+    def __init__(self, attentional=True):
+        super(DecoderBase, self).__init__()
+        self.attentional = attentional
+
+    @classmethod
+    def from_opt(cls, opt, embeddings):
+        """Alternate constructor.
+
+        Subclasses should override this method.
+        """
+
+        raise NotImplementedError

onmt_modules/decoder_transformer.py

@@ -0,0 +1,358 @@
+"""
+Implementation of "Attention is All You Need"
+"""
+
+import torch
+import torch.nn as nn
+
+from .decoder import DecoderBase
+from .multi_headed_attn import MultiHeadedAttention
+from .average_attn import AverageAttention
+from .position_ffn import PositionwiseFeedForward
+from .misc import sequence_mask
+
+
+class TransformerDecoderLayer(nn.Module):
+    """Transformer Decoder layer block in Pre-Norm style.
+    Pre-Norm style is an improvement w.r.t. Original paper's Post-Norm style,
+    providing better converge speed and performance. This is also the actual
+    implementation in tensor2tensor and also avalable in fairseq.
+    See https://tunz.kr/post/4 and :cite:`DeeperTransformer`.
+
+    .. mermaid::
+
+        graph LR
+        %% "*SubLayer" can be self-attn, src-attn or feed forward block
+            A(input) --> B[Norm]
+            B --> C["*SubLayer"]
+            C --> D[Drop]
+            D --> E((+))
+            A --> E
+            E --> F(out)
+
+
+    Args:
+        d_model (int): the dimension of keys/values/queries in
+            :class:`MultiHeadedAttention`, also the input size of
+            the first-layer of the :class:`PositionwiseFeedForward`.
+        heads (int): the number of heads for MultiHeadedAttention.
+        d_ff (int): the second-layer of the :class:`PositionwiseFeedForward`.
+        dropout (float): dropout in residual, self-attn(dot) and feed-forward
+        attention_dropout (float): dropout in context_attn (and self-attn(avg))
+        self_attn_type (string): type of self-attention scaled-dot, average
+        max_relative_positions (int):
+            Max distance between inputs in relative positions representations
+        aan_useffn (bool): Turn on the FFN layer in the AAN decoder
+        full_context_alignment (bool):
+            whether enable an extra full context decoder forward for alignment
+        alignment_heads (int):
+            N. of cross attention heads to use for alignment guiding
+    """
+
+    def __init__(self, d_model, heads, d_ff, dropout, attention_dropout,
+                 self_attn_type="scaled-dot", max_relative_positions=0,
+                 aan_useffn=False, full_context_alignment=False,
+                 alignment_heads=0):
+        super(TransformerDecoderLayer, self).__init__()
+
+        if self_attn_type == "scaled-dot":
+            self.self_attn = MultiHeadedAttention(
+                heads, d_model, dropout=attention_dropout,
+                max_relative_positions=max_relative_positions)
+        elif self_attn_type == "average":
+            self.self_attn = AverageAttention(d_model,
+                                              dropout=attention_dropout,
+                                              aan_useffn=aan_useffn)
+
+        self.context_attn = MultiHeadedAttention(
+            heads, d_model, dropout=attention_dropout)
+        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
+        self.layer_norm_1 = nn.LayerNorm(d_model, eps=1e-6)
+        self.layer_norm_2 = nn.LayerNorm(d_model, eps=1e-6)
+        self.drop = nn.Dropout(dropout)
+        self.full_context_alignment = full_context_alignment
+        self.alignment_heads = alignment_heads
+
+    def forward(self, *args, **kwargs):
+        """ Extend `_forward` for (possibly) multiple decoder pass:
+        Always a default (future masked) decoder forward pass,
+        Possibly a second future aware decoder pass for joint learn
+        full context alignement, :cite:`garg2019jointly`.
+
+        Args:
+            * All arguments of _forward.
+            with_align (bool): whether return alignment attention.
+
+        Returns:
+            (FloatTensor, FloatTensor, FloatTensor or None):
+
+            * output ``(batch_size, T, model_dim)``
+            * top_attn ``(batch_size, T, src_len)``
+            * attn_align ``(batch_size, T, src_len)`` or None
+        """
+        with_align = kwargs.pop('with_align', False)
+        output, attns = self._forward(*args, **kwargs)
+        top_attn = attns[:, 0, :, :].contiguous()
+        attn_align = None
+        if with_align:
+            if self.full_context_alignment:
+                # return _, (B, Q_len, K_len)
+                _, attns = self._forward(*args, **kwargs, future=True)
+
+            if self.alignment_heads > 0:
+                attns = attns[:, :self.alignment_heads, :, :].contiguous()
+            # layer average attention across heads, get ``(B, Q, K)``
+            # Case 1: no full_context, no align heads -> layer avg baseline
+            # Case 2: no full_context, 1 align heads -> guided align
+            # Case 3: full_context, 1 align heads -> full cte guided align
+            attn_align = attns.mean(dim=1)
+        return output, top_attn, attn_align
+
+    def _forward(self, inputs, memory_bank, src_pad_mask, tgt_pad_mask,
+                 layer_cache=None, step=None, future=False):
+        """ A naive forward pass for transformer decoder.
+
+        # T: could be 1 in the case of stepwise decoding or tgt_len
+
+        Args:
+            inputs (FloatTensor): ``(batch_size, T, model_dim)``
+            memory_bank (FloatTensor): ``(batch_size, src_len, model_dim)``
+            src_pad_mask (LongTensor): ``(batch_size, 1, src_len)``
+            tgt_pad_mask (LongTensor): ``(batch_size, 1, T)``
+            layer_cache (dict or None): cached layer info when stepwise decode
+            step (int or None): stepwise decoding counter
+            future (bool): If set True, do not apply future_mask.
+
+        Returns:
+            (FloatTensor, FloatTensor):
+
+            * output ``(batch_size, T, model_dim)``
+            * attns ``(batch_size, head, T, src_len)``
+
+        """
+        dec_mask = None
+
+        if step is None:
+            tgt_len = tgt_pad_mask.size(-1)
+            if not future:  # apply future_mask, result mask in (B, T, T)
+                future_mask = torch.ones(
+                    [tgt_len, tgt_len],
+                    device=tgt_pad_mask.device,
+                    dtype=torch.uint8)
+                future_mask = future_mask.triu_(1).view(1, tgt_len, tgt_len)
+                # BoolTensor was introduced in pytorch 1.2
+                try:
+                    future_mask = future_mask.bool()
+                except AttributeError:
+                    pass
+                dec_mask = torch.gt(tgt_pad_mask + future_mask, 0)
+            else:  # only mask padding, result mask in (B, 1, T)
+                dec_mask = tgt_pad_mask
+
+        input_norm = self.layer_norm_1(inputs)
+
+        if isinstance(self.self_attn, MultiHeadedAttention):
+            query, _ = self.self_attn(input_norm, input_norm, input_norm,
+                                      mask=dec_mask,
+                                      layer_cache=layer_cache,
+                                      attn_type="self")
+        elif isinstance(self.self_attn, AverageAttention):
+            query, _ = self.self_attn(input_norm, mask=dec_mask,
+                                      layer_cache=layer_cache, step=step)
+
+        query = self.drop(query) + inputs
+
+        query_norm = self.layer_norm_2(query)
+        mid, attns = self.context_attn(memory_bank, memory_bank, query_norm,
+                                       mask=src_pad_mask,
+                                       layer_cache=layer_cache,
+                                       attn_type="context")
+        output = self.feed_forward(self.drop(mid) + query)
+
+        return output, attns
+
+    def update_dropout(self, dropout, attention_dropout):
+        self.self_attn.update_dropout(attention_dropout)
+        self.context_attn.update_dropout(attention_dropout)
+        self.feed_forward.update_dropout(dropout)
+        self.drop.p = dropout
+
+
+class TransformerDecoder(DecoderBase):
+    """The Transformer decoder from "Attention is All You Need".
+    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`
+
+    .. mermaid::
+
+       graph BT
+          A[input]
+          B[multi-head self-attn]
+          BB[multi-head src-attn]
+          C[feed forward]
+          O[output]
+          A --> B
+          B --> BB
+          BB --> C
+          C --> O
+
+
+    Args:
+        num_layers (int): number of encoder layers.
+        d_model (int): size of the model
+        heads (int): number of heads
+        d_ff (int): size of the inner FF layer
+        copy_attn (bool): if using a separate copy attention
+        self_attn_type (str): type of self-attention scaled-dot, average
+        dropout (float): dropout in residual, self-attn(dot) and feed-forward
+        attention_dropout (float): dropout in context_attn (and self-attn(avg))
+        embeddings (onmt.modules.Embeddings):
+            embeddings to use, should have positional encodings
+        max_relative_positions (int):
+            Max distance between inputs in relative positions representations
+        aan_useffn (bool): Turn on the FFN layer in the AAN decoder
+        full_context_alignment (bool):
+            whether enable an extra full context decoder forward for alignment
+        alignment_layer (int): N° Layer to supervise with for alignment guiding
+        alignment_heads (int):
+            N. of cross attention heads to use for alignment guiding
+    """
+
+    def __init__(self, num_layers, d_model, heads, d_ff,
+                 copy_attn, self_attn_type, dropout, attention_dropout,
+                 embeddings, max_relative_positions, aan_useffn,
+                 full_context_alignment, alignment_layer,
+                 alignment_heads):
+        super(TransformerDecoder, self).__init__()
+
+        self.embeddings = embeddings
+
+        # Decoder State
+        self.state = {}
+
+        self.transformer_layers = nn.ModuleList(
+            [TransformerDecoderLayer(d_model, heads, d_ff, dropout,
+             attention_dropout, self_attn_type=self_attn_type,
+             max_relative_positions=max_relative_positions,
+             aan_useffn=aan_useffn,
+             full_context_alignment=full_context_alignment,
+             alignment_heads=alignment_heads)
+             for i in range(num_layers)])
+
+        # previously, there was a GlobalAttention module here for copy
+        # attention. But it was never actually used -- the "copy" attention
+        # just reuses the context attention.
+        self._copy = copy_attn
+        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
+
+        self.alignment_layer = alignment_layer
+
+    @classmethod
+    def from_opt(cls, opt, embeddings):
+        """Alternate constructor."""
+        return cls(
+            opt.dec_layers,
+            opt.dec_rnn_size,
+            opt.heads,
+            opt.transformer_ff,
+            opt.copy_attn,
+            opt.self_attn_type,
+            opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
+            opt.attention_dropout[0] if type(opt.attention_dropout)
+            is list else opt.dropout,
+            embeddings,
+            opt.max_relative_positions,
+            opt.aan_useffn,
+            opt.full_context_alignment,
+            opt.alignment_layer,
+            alignment_heads=opt.alignment_heads)
+
+    def init_state(self, src, memory_bank, enc_hidden):
+        """Initialize decoder state."""
+        self.state["src"] = src
+        self.state["cache"] = None
+
+    def map_state(self, fn):
+        def _recursive_map(struct, batch_dim=0):
+            for k, v in struct.items():
+                if v is not None:
+                    if isinstance(v, dict):
+                        _recursive_map(v)
+                    else:
+                        struct[k] = fn(v, batch_dim)
+
+        self.state["src"] = fn(self.state["src"], 1)
+        if self.state["cache"] is not None:
+            _recursive_map(self.state["cache"])
+
+    def detach_state(self):
+        self.state["src"] = self.state["src"].detach()
+
+    def forward(self, tgt, memory_bank, step=None, **kwargs):
+        """Decode, possibly stepwise."""
+        if step == 0:
+            self._init_cache(memory_bank)
+
+        tgt_words = tgt[:, :, 0].transpose(0, 1)
+
+        emb = self.embeddings(tgt, step=step)
+        assert emb.dim() == 3  # len x batch x embedding_dim
+
+        output = emb.transpose(0, 1).contiguous()
+        src_memory_bank = memory_bank.transpose(0, 1).contiguous()
+
+        pad_idx = self.embeddings.word_padding_idx
+        src_lens = kwargs["memory_lengths"]
+        src_max_len = self.state["src"].shape[0]
+        src_pad_mask = ~sequence_mask(src_lens, src_max_len).unsqueeze(1)
+        tgt_pad_mask = tgt_words.data.eq(pad_idx).unsqueeze(1)  # [B, 1, T_tgt]
+
+        with_align = kwargs.pop('with_align', False)
+        attn_aligns = []
+
+        for i, layer in enumerate(self.transformer_layers):
+            layer_cache = self.state["cache"]["layer_{}".format(i)] \
+                if step is not None else None
+            output, attn, attn_align = layer(
+                output,
+                src_memory_bank,
+                src_pad_mask,
+                tgt_pad_mask,
+                layer_cache=layer_cache,
+                step=step,
+                with_align=with_align)
+            if attn_align is not None:
+                attn_aligns.append(attn_align)
+
+        output = self.layer_norm(output)
+        dec_outs = output.transpose(0, 1).contiguous()
+        attn = attn.transpose(0, 1).contiguous()
+
+        attns = {"std": attn}
+        if self._copy:
+            attns["copy"] = attn
+        if with_align:
+            attns["align"] = attn_aligns[self.alignment_layer]  # `(B, Q, K)`
+            # attns["align"] = torch.stack(attn_aligns, 0).mean(0)  # All avg
+
+        # TODO change the way attns is returned dict => list or tuple (onnx)
+        return dec_outs, attns
+
+    def _init_cache(self, memory_bank):
+        self.state["cache"] = {}
+        batch_size = memory_bank.size(1)
+        depth = memory_bank.size(-1)
+
+        for i, layer in enumerate(self.transformer_layers):
+            layer_cache = {"memory_keys": None, "memory_values": None}
+            if isinstance(layer.self_attn, AverageAttention):
+                layer_cache["prev_g"] = torch.zeros((batch_size, 1, depth),
+                                                    device=memory_bank.device)
+            else:
+                layer_cache["self_keys"] = None
+                layer_cache["self_values"] = None
+            self.state["cache"]["layer_{}".format(i)] = layer_cache
+
+    def update_dropout(self, dropout, attention_dropout):
+        self.embeddings.update_dropout(dropout)
+        for layer in self.transformer_layers:
+            layer.update_dropout(dropout, attention_dropout)

onmt_modules/embeddings.py

@@ -0,0 +1,52 @@
+""" Embeddings module """
+import math
+import warnings
+
+import torch
+import torch.nn as nn
+
+
+class PositionalEncoding(nn.Module):
+    """Sinusoidal positional encoding for non-recurrent neural networks.
+
+    Implementation based on "Attention Is All You Need"
+    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`
+
+    Args:
+       dropout (float): dropout parameter
+       dim (int): embedding size
+    """
+
+    def __init__(self, dropout, dim, max_len=5000):
+        if dim % 2 != 0:
+            raise ValueError("Cannot use sin/cos positional encoding with "
+                             "odd dim (got dim={:d})".format(dim))
+        pe = torch.zeros(max_len, dim)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *
+                             -(math.log(10000.0) / dim)))
+        pe[:, 0::2] = torch.sin(position.float() * div_term)
+        pe[:, 1::2] = torch.cos(position.float() * div_term)
+        pe = pe.unsqueeze(1)
+        super(PositionalEncoding, self).__init__()
+        self.register_buffer('pe', pe)
+        self.dropout = nn.Dropout(p=dropout)
+        self.dim = dim
+
+    def forward(self, emb, step=None):
+        """Embed inputs.
+
+        Args:
+            emb (FloatTensor): Sequence of word vectors
+                ``(seq_len, batch_size, self.dim)``
+            step (int or NoneType): If stepwise (``seq_len = 1``), use
+                the encoding for this position.
+        """
+
+        emb = emb * math.sqrt(self.dim)
+        if step is None:
+            emb = emb + self.pe[:emb.size(0)]
+        else:
+            emb = emb + self.pe[step]
+        emb = self.dropout(emb)
+        return emb

onmt_modules/encoder.py

@@ -0,0 +1,58 @@
+"""Base class for encoders and generic multi encoders."""
+
+import torch.nn as nn
+
+from .misc import aeq
+
+
+class EncoderBase(nn.Module):
+    """
+    Base encoder class. Specifies the interface used by different encoder types
+    and required by :class:`onmt.Models.NMTModel`.
+
+    .. mermaid::
+
+       graph BT
+          A[Input]
+          subgraph RNN
+            C[Pos 1]
+            D[Pos 2]
+            E[Pos N]
+          end
+          F[Memory_Bank]
+          G[Final]
+          A-->C
+          A-->D
+          A-->E
+          C-->F
+          D-->F
+          E-->F
+          E-->G
+    """
+
+    @classmethod
+    def from_opt(cls, opt, embeddings=None):
+        raise NotImplementedError
+
+    def _check_args(self, src, lengths=None, hidden=None):
+        n_batch = src.size(1)
+        if lengths is not None:
+            n_batch_, = lengths.size()
+            aeq(n_batch, n_batch_)
+
+    def forward(self, src, lengths=None):
+        """
+        Args:
+            src (LongTensor):
+               padded sequences of sparse indices ``(src_len, batch, nfeat)``
+            lengths (LongTensor): length of each sequence ``(batch,)``
+
+
+        Returns:
+            (FloatTensor, FloatTensor):
+
+            * final encoder state, used to initialize decoder
+            * memory bank for attention, ``(src_len, batch, hidden)``
+        """
+
+        raise NotImplementedError

onmt_modules/encoder_transformer.py

@@ -0,0 +1,135 @@
+"""
+Implementation of "Attention is All You Need"
+"""
+
+import torch.nn as nn
+
+from .encoder import EncoderBase
+from .multi_headed_attn import MultiHeadedAttention
+from .position_ffn import PositionwiseFeedForward
+from .misc import sequence_mask
+
+
+class TransformerEncoderLayer(nn.Module):
+    """
+    A single layer of the transformer encoder.
+
+    Args:
+        d_model (int): the dimension of keys/values/queries in
+                   MultiHeadedAttention, also the input size of
+                   the first-layer of the PositionwiseFeedForward.
+        heads (int): the number of head for MultiHeadedAttention.
+        d_ff (int): the second-layer of the PositionwiseFeedForward.
+        dropout (float): dropout probability(0-1.0).
+    """
+
+    def __init__(self, d_model, heads, d_ff, dropout, attention_dropout,
+                 max_relative_positions=0):
+        super(TransformerEncoderLayer, self).__init__()
+
+        self.self_attn = MultiHeadedAttention(
+            heads, d_model, dropout=attention_dropout,
+            max_relative_positions=max_relative_positions)
+        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
+        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, inputs, mask):
+        """
+        Args:
+            inputs (FloatTensor): ``(batch_size, src_len, model_dim)``
+            mask (LongTensor): ``(batch_size, 1, src_len)``
+
+        Returns:
+            (FloatTensor):
+
+            * outputs ``(batch_size, src_len, model_dim)``
+        """
+        input_norm = self.layer_norm(inputs)
+        context, _ = self.self_attn(input_norm, input_norm, input_norm,
+                                    mask=mask, attn_type="self")
+        out = self.dropout(context) + inputs
+        return self.feed_forward(out)
+
+    def update_dropout(self, dropout, attention_dropout):
+        self.self_attn.update_dropout(attention_dropout)
+        self.feed_forward.update_dropout(dropout)
+        self.dropout.p = dropout
+
+
+class TransformerEncoder(EncoderBase):
+    """The Transformer encoder from "Attention is All You Need"
+    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`
+
+    .. mermaid::
+
+       graph BT
+          A[input]
+          B[multi-head self-attn]
+          C[feed forward]
+          O[output]
+          A --> B
+          B --> C
+          C --> O
+
+    Args:
+        num_layers (int): number of encoder layers
+        d_model (int): size of the model
+        heads (int): number of heads
+        d_ff (int): size of the inner FF layer
+        dropout (float): dropout parameters
+        embeddings (onmt.modules.Embeddings):
+          embeddings to use, should have positional encodings
+
+    Returns:
+        (torch.FloatTensor, torch.FloatTensor):
+
+        * embeddings ``(src_len, batch_size, model_dim)``
+        * memory_bank ``(src_len, batch_size, model_dim)``
+    """
+
+    def __init__(self, num_layers, d_model, heads, d_ff, dropout,
+                 attention_dropout, embeddings, max_relative_positions):
+        super(TransformerEncoder, self).__init__()
+
+        self.embeddings = embeddings
+        self.transformer = nn.ModuleList(
+            [TransformerEncoderLayer(
+                d_model, heads, d_ff, dropout, attention_dropout,
+                max_relative_positions=max_relative_positions)
+             for i in range(num_layers)])
+        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
+
+    @classmethod
+    def from_opt(cls, opt, embeddings):
+        """Alternate constructor."""
+        return cls(
+            opt.enc_layers,
+            opt.enc_rnn_size,
+            opt.heads,
+            opt.transformer_ff,
+            opt.dropout[0] if type(opt.dropout) is list else opt.dropout,
+            opt.attention_dropout[0] if type(opt.attention_dropout)
+            is list else opt.attention_dropout,
+            embeddings,
+            opt.max_relative_positions)
+
+    def forward(self, src, lengths=None):
+        """See :func:`EncoderBase.forward()`"""
+        self._check_args(src, lengths)
+
+        emb = self.embeddings(src)
+
+        out = emb.transpose(0, 1).contiguous()
+        mask = ~sequence_mask(lengths).unsqueeze(1)
+        # Run the forward pass of every layer of the tranformer.
+        for layer in self.transformer:
+            out = layer(out, mask)
+        out = self.layer_norm(out)
+
+        return emb, out.transpose(0, 1).contiguous(), lengths
+
+    def update_dropout(self, dropout, attention_dropout):
+        self.embeddings.update_dropout(dropout)
+        for layer in self.transformer:
+            layer.update_dropout(dropout, attention_dropout)

onmt_modules/misc.py

@@ -0,0 +1,173 @@
+# -*- coding: utf-8 -*-
+
+import torch
+import random
+import inspect
+from itertools import islice, repeat
+import os
+
+
+def split_corpus(path, shard_size, default=None):
+    """yield a `list` containing `shard_size` line of `path`,
+    or repeatly generate `default` if `path` is None.
+    """
+    if path is not None:
+        return _split_corpus(path, shard_size)
+    else:
+        return repeat(default)
+
+
+def _split_corpus(path, shard_size):
+    """Yield a `list` containing `shard_size` line of `path`.
+    """
+    with open(path, "rb") as f:
+        if shard_size <= 0:
+            yield f.readlines()
+        else:
+            while True:
+                shard = list(islice(f, shard_size))
+                if not shard:
+                    break
+                yield shard
+
+
+def aeq(*args):
+    """
+    Assert all arguments have the same value
+    """
+    arguments = (arg for arg in args)
+    first = next(arguments)
+    assert all(arg == first for arg in arguments), \
+        "Not all arguments have the same value: " + str(args)
+
+
+def sequence_mask(lengths, max_len=None):
+    """
+    Creates a boolean mask from sequence lengths.
+    """
+    batch_size = lengths.numel()
+    max_len = max_len or lengths.max()
+    return (torch.arange(0, max_len, device=lengths.device)
+            .type_as(lengths)
+            .repeat(batch_size, 1)
+            .lt(lengths.unsqueeze(1)))
+
+
+def tile(x, count, dim=0):
+    """
+    Tiles x on dimension dim count times.
+    """
+    perm = list(range(len(x.size())))
+    if dim != 0:
+        perm[0], perm[dim] = perm[dim], perm[0]
+        x = x.permute(perm).contiguous()
+    out_size = list(x.size())
+    out_size[0] *= count
+    batch = x.size(0)
+    x = x.view(batch, -1) \
+         .transpose(0, 1) \
+         .repeat(count, 1) \
+         .transpose(0, 1) \
+         .contiguous() \
+         .view(*out_size)
+    if dim != 0:
+        x = x.permute(perm).contiguous()
+    return x
+
+
+def use_gpu(opt):
+    """
+    Creates a boolean if gpu used
+    """
+    return (hasattr(opt, 'gpu_ranks') and len(opt.gpu_ranks) > 0) or \
+        (hasattr(opt, 'gpu') and opt.gpu > -1)
+
+
+def set_random_seed(seed, is_cuda):
+    """Sets the random seed."""
+    if seed > 0:
+        torch.manual_seed(seed)
+        # this one is needed for torchtext random call (shuffled iterator)
+        # in multi gpu it ensures datasets are read in the same order
+        random.seed(seed)
+        # some cudnn methods can be random even after fixing the seed
+        # unless you tell it to be deterministic
+        torch.backends.cudnn.deterministic = True
+
+    if is_cuda and seed > 0:
+        # These ensure same initialization in multi gpu mode
+        torch.cuda.manual_seed(seed)
+
+
+def generate_relative_positions_matrix(length, max_relative_positions,
+                                       cache=False):
+    """Generate the clipped relative positions matrix
+       for a given length and maximum relative positions"""
+    if cache:
+        distance_mat = torch.arange(-length+1, 1, 1).unsqueeze(0)
+    else:
+        range_vec = torch.arange(length)
+        range_mat = range_vec.unsqueeze(-1).expand(-1, length).transpose(0, 1)
+        distance_mat = range_mat - range_mat.transpose(0, 1)
+    distance_mat_clipped = torch.clamp(distance_mat,
+                                       min=-max_relative_positions,
+                                       max=max_relative_positions)
+    # Shift values to be >= 0
+    final_mat = distance_mat_clipped + max_relative_positions
+    return final_mat
+
+
+def relative_matmul(x, z, transpose):
+    """Helper function for relative positions attention."""
+    batch_size = x.shape[0]
+    heads = x.shape[1]
+    length = x.shape[2]
+    x_t = x.permute(2, 0, 1, 3)
+    x_t_r = x_t.reshape(length, heads * batch_size, -1)
+    if transpose:
+        z_t = z.transpose(1, 2)
+        x_tz_matmul = torch.matmul(x_t_r, z_t)
+    else:
+        x_tz_matmul = torch.matmul(x_t_r, z)
+    x_tz_matmul_r = x_tz_matmul.reshape(length, batch_size, heads, -1)
+    x_tz_matmul_r_t = x_tz_matmul_r.permute(1, 2, 0, 3)
+    return x_tz_matmul_r_t
+
+
+def fn_args(fun):
+    """Returns the list of function arguments name."""
+    return inspect.getfullargspec(fun).args
+
+
+def report_matrix(row_label, column_label, matrix):
+    header_format = "{:>10.10} " + "{:>10.7} " * len(row_label)
+    row_format = "{:>10.10} " + "{:>10.7f} " * len(row_label)
+    output = header_format.format("", *row_label) + '\n'
+    for word, row in zip(column_label, matrix):
+        max_index = row.index(max(row))
+        row_format = row_format.replace(
+            "{:>10.7f} ", "{:*>10.7f} ", max_index + 1)
+        row_format = row_format.replace(
+            "{:*>10.7f} ", "{:>10.7f} ", max_index)
+        output += row_format.format(word, *row) + '\n'
+        row_format = "{:>10.10} " + "{:>10.7f} " * len(row_label)
+    return output
+
+
+def check_model_config(model_config, root):
+    # we need to check the model path + any tokenizer path
+    for model in model_config["models"]:
+        model_path = os.path.join(root, model)
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(
+                "{} from model {} does not exist".format(
+                    model_path, model_config["id"]))
+    if "tokenizer" in model_config.keys():
+        if "params" in model_config["tokenizer"].keys():
+            for k, v in model_config["tokenizer"]["params"].items():
+                if k.endswith("path"):
+                    tok_path = os.path.join(root, v)
+                    if not os.path.exists(tok_path):
+                        raise FileNotFoundError(
+                            "{} from model {} does not exist".format(
+                                tok_path, model_config["id"]))

onmt_modules/multi_headed_attn.py

@@ -0,0 +1,230 @@
+""" Multi-Head Attention module """
+import math
+import torch
+import torch.nn as nn
+
+from .misc import generate_relative_positions_matrix,\
+                            relative_matmul
+# from onmt.utils.misc import aeq
+
+
+class MultiHeadedAttention(nn.Module):
+    """Multi-Head Attention module from "Attention is All You Need"
+    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.
+
+    Similar to standard `dot` attention but uses
+    multiple attention distributions simulataneously
+    to select relevant items.
+
+    .. mermaid::
+
+       graph BT
+          A[key]
+          B[value]
+          C[query]
+          O[output]
+          subgraph Attn
+            D[Attn 1]
+            E[Attn 2]
+            F[Attn N]
+          end
+          A --> D
+          C --> D
+          A --> E
+          C --> E
+          A --> F
+          C --> F
+          D --> O
+          E --> O
+          F --> O
+          B --> O
+
+    Also includes several additional tricks.
+
+    Args:
+       head_count (int): number of parallel heads
+       model_dim (int): the dimension of keys/values/queries,
+           must be divisible by head_count
+       dropout (float): dropout parameter
+    """
+
+    def __init__(self, head_count, model_dim, dropout=0.1,
+                 max_relative_positions=0):
+        assert model_dim % head_count == 0
+        self.dim_per_head = model_dim // head_count
+        self.model_dim = model_dim
+
+        super(MultiHeadedAttention, self).__init__()
+        self.head_count = head_count
+
+        self.linear_keys = nn.Linear(model_dim,
+                                     head_count * self.dim_per_head)
+        self.linear_values = nn.Linear(model_dim,
+                                       head_count * self.dim_per_head)
+        self.linear_query = nn.Linear(model_dim,
+                                      head_count * self.dim_per_head)
+        self.softmax = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(dropout)
+        self.final_linear = nn.Linear(model_dim, model_dim)
+
+        self.max_relative_positions = max_relative_positions
+
+        if max_relative_positions > 0:
+            vocab_size = max_relative_positions * 2 + 1
+            self.relative_positions_embeddings = nn.Embedding(
+                vocab_size, self.dim_per_head)
+
+    def forward(self, key, value, query, mask=None,
+                layer_cache=None, attn_type=None):
+        """
+        Compute the context vector and the attention vectors.
+
+        Args:
+           key (FloatTensor): set of `key_len`
+               key vectors ``(batch, key_len, dim)``
+           value (FloatTensor): set of `key_len`
+               value vectors ``(batch, key_len, dim)``
+           query (FloatTensor): set of `query_len`
+               query vectors  ``(batch, query_len, dim)``
+           mask: binary mask 1/0 indicating which keys have
+               zero / non-zero attention ``(batch, query_len, key_len)``
+        Returns:
+           (FloatTensor, FloatTensor):
+
+           * output context vectors ``(batch, query_len, dim)``
+           * Attention vector in heads ``(batch, head, query_len, key_len)``.
+        """
+
+        # CHECKS
+        # batch, k_len, d = key.size()
+        # batch_, k_len_, d_ = value.size()
+        # aeq(batch, batch_)
+        # aeq(k_len, k_len_)
+        # aeq(d, d_)
+        # batch_, q_len, d_ = query.size()
+        # aeq(batch, batch_)
+        # aeq(d, d_)
+        # aeq(self.model_dim % 8, 0)
+        # if mask is not None:
+        #    batch_, q_len_, k_len_ = mask.size()
+        #    aeq(batch_, batch)
+        #    aeq(k_len_, k_len)
+        #    aeq(q_len_ == q_len)
+        # END CHECKS
+
+        batch_size = key.size(0)
+        dim_per_head = self.dim_per_head
+        head_count = self.head_count
+        key_len = key.size(1)
+        query_len = query.size(1)
+
+        def shape(x):
+            """Projection."""
+            return x.view(batch_size, -1, head_count, dim_per_head) \
+                .transpose(1, 2)
+
+        def unshape(x):
+            """Compute context."""
+            return x.transpose(1, 2).contiguous() \
+                    .view(batch_size, -1, head_count * dim_per_head)
+
+        # 1) Project key, value, and query.
+        if layer_cache is not None:
+            if attn_type == "self":
+                query, key, value = self.linear_query(query),\
+                                    self.linear_keys(query),\
+                                    self.linear_values(query)
+                key = shape(key)
+                value = shape(value)
+                if layer_cache["self_keys"] is not None:
+                    key = torch.cat(
+                        (layer_cache["self_keys"], key),
+                        dim=2)
+                if layer_cache["self_values"] is not None:
+                    value = torch.cat(
+                        (layer_cache["self_values"], value),
+                        dim=2)
+                layer_cache["self_keys"] = key
+                layer_cache["self_values"] = value
+            elif attn_type == "context":
+                query = self.linear_query(query)
+                if layer_cache["memory_keys"] is None:
+                    key, value = self.linear_keys(key),\
+                                 self.linear_values(value)
+                    key = shape(key)
+                    value = shape(value)
+                else:
+                    key, value = layer_cache["memory_keys"],\
+                               layer_cache["memory_values"]
+                layer_cache["memory_keys"] = key
+                layer_cache["memory_values"] = value
+        else:
+            key = self.linear_keys(key)
+            value = self.linear_values(value)
+            query = self.linear_query(query)
+            key = shape(key)
+            value = shape(value)
+
+        if self.max_relative_positions > 0 and attn_type == "self":
+            key_len = key.size(2)
+            # 1 or key_len x key_len
+            relative_positions_matrix = generate_relative_positions_matrix(
+                key_len, self.max_relative_positions,
+                cache=True if layer_cache is not None else False)
+            #  1 or key_len x key_len x dim_per_head
+            relations_keys = self.relative_positions_embeddings(
+                relative_positions_matrix.to(key.device))
+            #  1 or key_len x key_len x dim_per_head
+            relations_values = self.relative_positions_embeddings(
+                relative_positions_matrix.to(key.device))
+
+        query = shape(query)
+
+        key_len = key.size(2)
+        query_len = query.size(2)
+
+        # 2) Calculate and scale scores.
+        query = query / math.sqrt(dim_per_head)
+        # batch x num_heads x query_len x key_len
+        query_key = torch.matmul(query, key.transpose(2, 3))
+
+        if self.max_relative_positions > 0 and attn_type == "self":
+            scores = query_key + relative_matmul(query, relations_keys, True)
+        else:
+            scores = query_key
+        scores = scores.float()
+
+        if mask is not None:
+            mask = mask.unsqueeze(1)  # [B, 1, 1, T_values]
+            scores = scores.masked_fill(mask, -1e18)
+
+        # 3) Apply attention dropout and compute context vectors.
+        attn = self.softmax(scores).to(query.dtype)
+        drop_attn = self.dropout(attn)
+
+        context_original = torch.matmul(drop_attn, value)
+
+        if self.max_relative_positions > 0 and attn_type == "self":
+            context = unshape(context_original
+                              + relative_matmul(drop_attn,
+                                                relations_values,
+                                                False))
+        else:
+            context = unshape(context_original)
+
+        output = self.final_linear(context)
+        # CHECK
+        # batch_, q_len_, d_ = output.size()
+        # aeq(q_len, q_len_)
+        # aeq(batch, batch_)
+        # aeq(d, d_)
+
+        # Return multi-head attn
+        attns = attn \
+            .view(batch_size, head_count,
+                  query_len, key_len)
+
+        return output, attns
+
+    def update_dropout(self, dropout):
+        self.dropout.p = dropout

onmt_modules/position_ffn.py

@@ -0,0 +1,41 @@
+"""Position feed-forward network from "Attention is All You Need"."""
+
+import torch.nn as nn
+
+
+class PositionwiseFeedForward(nn.Module):
+    """ A two-layer Feed-Forward-Network with residual layer norm.
+
+    Args:
+        d_model (int): the size of input for the first-layer of the FFN.
+        d_ff (int): the hidden layer size of the second-layer
+            of the FNN.
+        dropout (float): dropout probability in :math:`[0, 1)`.
+    """
+
+    def __init__(self, d_model, d_ff, dropout=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.w_1 = nn.Linear(d_model, d_ff)
+        self.w_2 = nn.Linear(d_ff, d_model)
+        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
+        self.dropout_1 = nn.Dropout(dropout)
+        self.relu = nn.ReLU()
+        self.dropout_2 = nn.Dropout(dropout)
+
+    def forward(self, x):
+        """Layer definition.
+
+        Args:
+            x: ``(batch_size, input_len, model_dim)``
+
+        Returns:
+            (FloatTensor): Output ``(batch_size, input_len, model_dim)``.
+        """
+
+        inter = self.dropout_1(self.relu(self.w_1(self.layer_norm(x))))
+        output = self.dropout_2(self.w_2(inter))
+        return output + x
+
+    def update_dropout(self, dropout):
+        self.dropout_1.p = dropout
+        self.dropout_2.p = dropout

override_decoder.py

@@ -0,0 +1,63 @@
+from onmt_modules.decoder_transformer import TransformerDecoder
+from onmt_modules.misc import sequence_mask
+
+  
+class OnmtDecoder_1(TransformerDecoder):
+    # overide forward
+    # without teacher forcing for stop
+    def forward(self, tgt, memory_bank, step=None, **kwargs):
+        """Decode, possibly stepwise."""
+        if step == 0:
+            self._init_cache(memory_bank)
+        
+        if step is None:
+            tgt_lens = kwargs["tgt_lengths"]
+        else:    
+            tgt_words = kwargs["tgt_words"]
+        
+        emb = self.embeddings(tgt, step=step)
+        assert emb.dim() == 3  # len x batch x embedding_dim
+
+        output = emb.transpose(0, 1).contiguous()
+        src_memory_bank = memory_bank.transpose(0, 1).contiguous()
+
+        pad_idx = self.embeddings.word_padding_idx
+        src_lens = kwargs["memory_lengths"]
+        src_max_len = self.state["src"].shape[0]
+        src_pad_mask = ~sequence_mask(src_lens, src_max_len).unsqueeze(1)
+        if step is None:
+            tgt_max_len = tgt_lens.max()
+            tgt_pad_mask = ~sequence_mask(tgt_lens, tgt_max_len).unsqueeze(1)
+        else:    
+            tgt_pad_mask = tgt_words.data.eq(pad_idx).unsqueeze(1)
+
+        with_align = kwargs.pop('with_align', False)
+        attn_aligns = []
+
+        for i, layer in enumerate(self.transformer_layers):
+            layer_cache = self.state["cache"]["layer_{}".format(i)] \
+                if step is not None else None
+            output, attn, attn_align = layer(
+                output,
+                src_memory_bank,
+                src_pad_mask,
+                tgt_pad_mask,
+                layer_cache=layer_cache,
+                step=step,
+                with_align=with_align)
+            if attn_align is not None:
+                attn_aligns.append(attn_align)
+
+        output = self.layer_norm(output)
+        dec_outs = output.transpose(0, 1).contiguous()
+        attn = attn.transpose(0, 1).contiguous()
+
+        attns = {"std": attn}
+        if self._copy:
+            attns["copy"] = attn
+        if with_align:
+            attns["align"] = attn_aligns[self.alignment_layer]  # `(B, Q, K)`
+            # attns["align"] = torch.stack(attn_aligns, 0).mean(0)  # All avg
+
+        # TODO change the way attns is returned dict => list or tuple (onnx)
+        return dec_outs, attns

prepare_train_data.py

@@ -0,0 +1,90 @@
+import os
+import pickle
+import numpy as np
+import scipy.fftpack
+import soundfile as sf
+from utils import pySTFT
+from scipy import signal
+from librosa.filters import mel
+from utils import butter_highpass
+
+import torch
+import torch.nn.functional as F
+from model_sea import Generator as Model 
+from hparams_sea import hparams
+
+    
+mel_basis = mel(16000, 1024, fmin=90, fmax=7600, n_mels=80).T
+min_level = np.exp(-100 / 20 * np.log(10))
+b, a = butter_highpass(30, 16000, order=5)
+
+mfcc_mean, mfcc_std, dctmx = pickle.load(open('assets/mfcc_stats.pkl', 'rb'))
+spk2emb = pickle.load(open('assets/spk2emb_82.pkl', 'rb'))
+
+rootDir = "assets/vctk16-train-wav"
+targetDir_sp = 'assets/vctk16-train-sp-mel'
+targetDir_cep = 'assets/vctk16-train-cep-mel'
+targetDir_cd = 'assets/vctk16-train-teacher'
+
+device = 'cuda:0'
+
+G = Model(hparams).eval().to(device)
+            
+g_checkpoint = torch.load('assets/sea.ckpt', map_location=lambda storage, loc: storage)
+G.load_state_dict(g_checkpoint['model'], strict=True)
+
+
+metadata = []
+dirName, subdirList, _ = next(os.walk(rootDir))
+
+for subdir in sorted(subdirList):
+    print(subdir)
+    
+    if not os.path.exists(os.path.join(targetDir_sp, subdir)):
+        os.makedirs(os.path.join(targetDir_sp, subdir))
+    if not os.path.exists(os.path.join(targetDir_cep, subdir)):
+        os.makedirs(os.path.join(targetDir_cep, subdir)) 
+    if not os.path.exists(os.path.join(targetDir_cd, subdir)):
+        os.makedirs(os.path.join(targetDir_cd, subdir))     
+    
+    submeta = []
+    submeta.append(subdir)
+    submeta.append(spk2emb[subdir])
+    
+    _,_, fileList = next(os.walk(os.path.join(dirName,subdir)))
+    
+    for fileName in sorted(fileList):
+        x, fs = sf.read(os.path.join(dirName,subdir,fileName))
+        if x.shape[0] % 256 == 0:
+            x = np.concatenate((x, np.array([1e-06])), axis=0)
+        y = signal.filtfilt(b, a, x)
+        D = pySTFT(y * 0.96).T
+        D_mel = np.dot(D, mel_basis)
+        D_db = 20 * np.log10(np.maximum(min_level, D_mel))
+
+        # mel sp
+        S = (D_db + 80) / 100 
+                
+        # mel cep
+        cc_tmp = S.dot(dctmx)
+        cc_norm = (cc_tmp - mfcc_mean) / mfcc_std
+        S = np.clip(S, 0, 1)
+        
+        # teacher code
+        cc_torch = torch.from_numpy(cc_norm[:,0:20].astype(np.float32)).unsqueeze(0).to(device)
+        with torch.no_grad():
+            codes = G.encode(cc_torch, torch.ones_like(cc_torch[:,:,0])).squeeze(0)
+            
+        np.save(os.path.join(targetDir_cd, subdir, fileName[:-4]), 
+                codes.cpu().numpy(), allow_pickle=False)
+        np.save(os.path.join(targetDir_sp, subdir, fileName[:-4]), 
+                S.astype(np.float32), allow_pickle=False)
+        np.save(os.path.join(targetDir_cep, subdir, fileName[:-4]), 
+                cc_norm.astype(np.float32), allow_pickle=False)
+        
+        submeta.append(subdir+'/'+fileName[:-4]+'.npy')
+            
+    metadata.append(submeta)
+        
+with open('./assets/train_vctk.meta', 'wb') as handle:
+    pickle.dump(metadata, handle)  

solver_1.py

@@ -0,0 +1,146 @@
+import os
+import time
+import pickle
+import datetime
+import itertools
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from onmt_modules.misc import sequence_mask
+from model_autopst import Generator_1 as Predictor
+
+
+
+class Solver(object):
+
+    def __init__(self, data_loader, config, hparams):
+        """Initialize configurations."""
+
+        
+        self.data_loader = data_loader
+        self.hparams = hparams
+        self.gate_threshold = hparams.gate_threshold
+        
+        self.use_cuda = torch.cuda.is_available()
+        self.device = torch.device('cuda:{}'.format(config.device_id) if self.use_cuda else 'cpu')
+        self.num_iters = config.num_iters
+        self.log_step = config.log_step
+        
+        # Build the model
+        self.build_model()
+    
+            
+    def build_model(self):
+        
+        self.P = Predictor(self.hparams)
+        
+        self.optimizer = torch.optim.Adam(self.P.parameters(), 0.0001, [0.9, 0.999])
+        
+        self.P.to(self.device)
+        
+        self.BCELoss = torch.nn.BCEWithLogitsLoss().to(self.device)    
+    
+                
+    def train(self):
+        # Set data loader
+        data_loader = self.data_loader
+        data_iter = iter(data_loader)
+        
+        
+        # Print logs in specified order
+        keys = ['P/loss_tx2sp', 'P/loss_stop_sp']
+        
+            
+        # Start training.
+        print('Start training...')
+        start_time = time.time()
+        for i in range(self.num_iters):
+
+            try:
+                sp_real, cep_real, cd_real, _, num_rep_sync, len_real, _, len_short_sync, spk_emb = next(data_iter)
+            except:
+                data_iter = iter(data_loader)
+                sp_real, cep_real, cd_real, _, num_rep_sync, len_real, _, len_short_sync, spk_emb = next(data_iter)
+                
+            
+            sp_real = sp_real.to(self.device)
+            cep_real = cep_real.to(self.device)
+            cd_real = cd_real.to(self.device)
+            len_real = len_real.to(self.device)
+            spk_emb = spk_emb.to(self.device)
+            num_rep_sync = num_rep_sync.to(self.device)
+            len_short_sync = len_short_sync.to(self.device)
+            
+            
+            # real spect masks
+            mask_sp_real = ~sequence_mask(len_real, sp_real.size(1))
+            mask_long = (~mask_sp_real).float()
+            
+            len_real_mask = torch.min(len_real + 10, 
+                                      torch.full_like(len_real, sp_real.size(1)))
+            loss_tx2sp_mask = sequence_mask(len_real_mask, sp_real.size(1)).float().unsqueeze(-1)
+            
+            # text input masks
+            codes_mask = sequence_mask(len_short_sync, num_rep_sync.size(1)).float()
+            
+            
+            # =================================================================================== #
+            #                                    2. Train                                         #
+            # =================================================================================== #
+            
+            self.P = self.P.train()
+            
+            
+            sp_real_sft = torch.zeros_like(sp_real)
+            sp_real_sft[:, 1:, :] = sp_real[:, :-1, :]    
+            
+            
+            spect_pred, stop_pred_sp = self.P(cep_real.transpose(2,1),
+                                              mask_long,
+                                              codes_mask,
+                                              num_rep_sync,
+                                              len_short_sync+1,
+                                              sp_real_sft.transpose(1,0), 
+                                              len_real+1,
+                                              spk_emb)
+                        
+            
+            loss_tx2sp = (F.mse_loss(spect_pred.permute(1,0,2), sp_real, reduction='none')
+                          * loss_tx2sp_mask).sum() / loss_tx2sp_mask.sum()
+                          
+            loss_stop_sp = self.BCELoss(stop_pred_sp.squeeze(-1).t(), mask_sp_real.float())
+            
+            loss_total = loss_tx2sp + loss_stop_sp
+            
+            # Backward and optimize
+            self.optimizer.zero_grad()
+            loss_total.backward()
+            self.optimizer.step()
+            
+
+            # Logging
+            loss = {}
+            loss['P/loss_tx2sp'] = loss_tx2sp.item()
+            loss['P/loss_stop_sp'] = loss_stop_sp.item()
+            
+
+            # =================================================================================== #
+            #                                 4. Miscellaneous                                    #
+            # =================================================================================== #
+
+            # Print out training information
+            if (i+1) % self.log_step == 0:
+                et = time.time() - start_time
+                et = str(datetime.timedelta(seconds=et))[:-7]
+                log = "Elapsed [{}], Iteration [{}/{}]".format(et, i+1, self.num_iters)
+                for tag in keys:
+                    log += ", {}: {:.8f}".format(tag, loss[tag])
+                print(log)
+                
+                
+            # Save model checkpoints.
+            if (i+1) % 10000 == 0:
+                torch.save({'model': self.P.state_dict(),
+                            'optimizer': self.optimizer.state_dict()}, f'./assets/{i+1}-A.ckpt')
+                print('Saved model checkpoints into assets ...')  

solver_2.py

@@ -0,0 +1,160 @@
+import os
+import time
+import pickle
+import datetime
+import itertools
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from onmt_modules.misc import sequence_mask
+from model_autopst import Generator_2 as Predictor
+
+
+
+class Solver(object):
+
+    def __init__(self, data_loader, config, hparams):
+        """Initialize configurations."""
+
+        
+        self.data_loader = data_loader
+        self.hparams = hparams
+        self.gate_threshold = hparams.gate_threshold
+        
+        self.use_cuda = torch.cuda.is_available()
+        self.device = torch.device('cuda:{}'.format(config.device_id) if self.use_cuda else 'cpu')
+        self.num_iters = config.num_iters
+        self.log_step = config.log_step
+        
+        # Build the model
+        self.build_model()
+    
+            
+    def build_model(self):
+        
+        self.P = Predictor(self.hparams)
+        self.freeze_layers(self.P.encoder_cd)
+        
+        self.optimizer = torch.optim.Adam(self.P.parameters(), 0.0001, [0.9, 0.999])
+        
+        self.P.to(self.device)
+        
+        self.BCELoss = torch.nn.BCEWithLogitsLoss().to(self.device)  
+        
+        
+        checkpoint = torch.load(self.hparams.pretrained_path, 
+                                map_location=lambda storage, loc: storage)
+        
+        self.P.load_state_dict(checkpoint['model'], strict=True)
+        print('Loaded pretrained encoder .........................................')
+        
+        
+    def freeze_layers(self, layer):
+        print('Fixing layers!')
+        for param in layer.parameters():
+            param.requires_grad = False
+    
+                
+    def train(self):
+        # Set data loader
+        data_loader = self.data_loader
+        data_iter = iter(data_loader)
+        
+        
+        # Print logs in specified order
+        keys = ['P/loss_tx2sp', 'P/loss_stop_sp']
+        
+            
+        # Start training.
+        print('Start training...')
+        start_time = time.time()
+        for i in range(self.num_iters):
+
+            try:
+                sp_real, cep_real, cd_real, num_rep, _, len_real, len_short, _, spk_emb = next(data_iter)
+            except:
+                data_iter = iter(data_loader)
+                sp_real, cep_real, cd_real, num_rep, _, len_real, len_short, _, spk_emb = next(data_iter)
+                
+            
+            sp_real = sp_real.to(self.device)
+            cep_real = cep_real.to(self.device)
+            cd_real = cd_real.to(self.device)
+            len_real = len_real.to(self.device)
+            spk_emb = spk_emb.to(self.device)
+            num_rep = num_rep.to(self.device)
+            len_short = len_short.to(self.device)
+            
+            
+            # real spect masks
+            mask_sp_real = ~sequence_mask(len_real, sp_real.size(1))
+            mask_long = (~mask_sp_real).float()
+            
+            len_real_mask = torch.min(len_real + 10, 
+                                      torch.full_like(len_real, sp_real.size(1)))
+            loss_tx2sp_mask = sequence_mask(len_real_mask, sp_real.size(1)).float().unsqueeze(-1)
+            
+            # text input masks
+            codes_mask = sequence_mask(len_short, num_rep.size(1)).float()
+            
+            
+            # =================================================================================== #
+            #                                    2. Train                                         #
+            # =================================================================================== #
+            
+            self.P = self.P.train()
+            
+            
+            sp_real_sft = torch.zeros_like(sp_real)
+            sp_real_sft[:, 1:, :] = sp_real[:, :-1, :]    
+            
+            
+            spect_pred, stop_pred_sp = self.P(cep_real.transpose(2,1),
+                                              mask_long,
+                                              codes_mask,
+                                              num_rep,
+                                              len_short+1,
+                                              sp_real_sft.transpose(1,0), 
+                                              len_real+1,
+                                              spk_emb)
+                        
+            
+            loss_tx2sp = (F.mse_loss(spect_pred.permute(1,0,2), sp_real, reduction='none')
+                          * loss_tx2sp_mask).sum() / loss_tx2sp_mask.sum()
+                          
+            loss_stop_sp = self.BCELoss(stop_pred_sp.squeeze(-1).t(), mask_sp_real.float())
+            
+            loss_total = loss_tx2sp + loss_stop_sp
+            
+            # Backward and optimize
+            self.optimizer.zero_grad()
+            loss_total.backward()
+            self.optimizer.step()
+            
+
+            # Logging
+            loss = {}
+            loss['P/loss_tx2sp'] = loss_tx2sp.item()
+            loss['P/loss_stop_sp'] = loss_stop_sp.item()
+            
+
+            # =================================================================================== #
+            #                                 4. Miscellaneous                                    #
+            # =================================================================================== #
+
+            # Print out training information
+            if (i+1) % self.log_step == 0:
+                et = time.time() - start_time
+                et = str(datetime.timedelta(seconds=et))[:-7]
+                log = "Elapsed [{}], Iteration [{}/{}]".format(et, i+1, self.num_iters)
+                for tag in keys:
+                    log += ", {}: {:.8f}".format(tag, loss[tag])
+                print(log)
+                
+                
+            # Save model checkpoints.
+            if (i+1) % 10000 == 0:
+                torch.save({'model': self.P.state_dict(),
+                            'optimizer': self.optimizer.state_dict()}, f'./assets/{i+1}-B.ckpt')
+                print('Saved model checkpoints into assets ...')  

tfcompat/hparam.py

@@ -0,0 +1,726 @@
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Hyperparameter values."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+import numbers
+import re
+
+import six
+
+## from tensorflow.contrib.training.python.training import hparam_pb2
+## from tensorflow.python.framework import ops
+## from tensorflow.python.util import compat
+## from tensorflow.python.util import deprecation
+
+# Define the regular expression for parsing a single clause of the input
+# (delimited by commas).  A legal clause looks like:
+#   <variable name>[<index>]? = <rhs>
+# where <rhs> is either a single token or [] enclosed list of tokens.
+# For example:  "var[1] = a" or "x = [1,2,3]"
+PARAM_RE = re.compile(r"""
+  (?P<name>[a-zA-Z][\w\.]*)      # variable name: "var" or "x"
+  (\[\s*(?P<index>\d+)\s*\])?  # (optional) index: "1" or None
+  \s*=\s*
+  ((?P<val>[^,\[]*)            # single value: "a" or None
+   |
+   \[(?P<vals>[^\]]*)\])       # list of values: None or "1,2,3"
+  ($|,\s*)""", re.VERBOSE)
+
+
+def _parse_fail(name, var_type, value, values):
+  """Helper function for raising a value error for bad assignment."""
+  raise ValueError(
+      'Could not parse hparam \'%s\' of type \'%s\' with value \'%s\' in %s' %
+      (name, var_type.__name__, value, values))
+
+
+def _reuse_fail(name, values):
+  """Helper function for raising a value error for reuse of name."""
+  raise ValueError('Multiple assignments to variable \'%s\' in %s' % (name,
+                                                                      values))
+
+
+def _process_scalar_value(name, parse_fn, var_type, m_dict, values,
+                          results_dictionary):
+  """Update results_dictionary with a scalar value.
+
+  Used to update the results_dictionary to be returned by parse_values when
+  encountering a clause with a scalar RHS (e.g.  "s=5" or "arr[0]=5".)
+
+  Mutates results_dictionary.
+
+  Args:
+    name: Name of variable in assignment ("s" or "arr").
+    parse_fn: Function for parsing the actual value.
+    var_type: Type of named variable.
+    m_dict: Dictionary constructed from regex parsing.
+      m_dict['val']: RHS value (scalar)
+      m_dict['index']: List index value (or None)
+    values: Full expression being parsed
+    results_dictionary: The dictionary being updated for return by the parsing
+      function.
+
+  Raises:
+    ValueError: If the name has already been used.
+  """
+  try:
+    parsed_value = parse_fn(m_dict['val'])
+  except ValueError:
+    _parse_fail(name, var_type, m_dict['val'], values)
+
+  # If no index is provided
+  if not m_dict['index']:
+    if name in results_dictionary:
+      _reuse_fail(name, values)
+    results_dictionary[name] = parsed_value
+  else:
+    if name in results_dictionary:
+      # The name has already been used as a scalar, then it
+      # will be in this dictionary and map to a non-dictionary.
+      if not isinstance(results_dictionary.get(name), dict):
+        _reuse_fail(name, values)
+    else:
+      results_dictionary[name] = {}
+
+    index = int(m_dict['index'])
+    # Make sure the index position hasn't already been assigned a value.
+    if index in results_dictionary[name]:
+      _reuse_fail('{}[{}]'.format(name, index), values)
+    results_dictionary[name][index] = parsed_value
+
+
+def _process_list_value(name, parse_fn, var_type, m_dict, values,
+                        results_dictionary):
+  """Update results_dictionary from a list of values.
+
+  Used to update results_dictionary to be returned by parse_values when
+  encountering a clause with a list RHS (e.g.  "arr=[1,2,3]".)
+
+  Mutates results_dictionary.
+
+  Args:
+    name: Name of variable in assignment ("arr").
+    parse_fn: Function for parsing individual values.
+    var_type: Type of named variable.
+    m_dict: Dictionary constructed from regex parsing.
+      m_dict['val']: RHS value (scalar)
+    values: Full expression being parsed
+    results_dictionary: The dictionary being updated for return by the parsing
+      function.
+
+  Raises:
+    ValueError: If the name has an index or the values cannot be parsed.
+  """
+  if m_dict['index'] is not None:
+    raise ValueError('Assignment of a list to a list index.')
+  elements = filter(None, re.split('[ ,]', m_dict['vals']))
+  # Make sure the name hasn't already been assigned a value
+  if name in results_dictionary:
+    raise _reuse_fail(name, values)
+  try:
+    results_dictionary[name] = [parse_fn(e) for e in elements]
+  except ValueError:
+    _parse_fail(name, var_type, m_dict['vals'], values)
+
+
+def _cast_to_type_if_compatible(name, param_type, value):
+  """Cast hparam to the provided type, if compatible.
+
+  Args:
+    name: Name of the hparam to be cast.
+    param_type: The type of the hparam.
+    value: The value to be cast, if compatible.
+
+  Returns:
+    The result of casting `value` to `param_type`.
+
+  Raises:
+    ValueError: If the type of `value` is not compatible with param_type.
+      * If `param_type` is a string type, but `value` is not.
+      * If `param_type` is a boolean, but `value` is not, or vice versa.
+      * If `param_type` is an integer type, but `value` is not.
+      * If `param_type` is a float type, but `value` is not a numeric type.
+  """
+  fail_msg = (
+      "Could not cast hparam '%s' of type '%s' from value %r" %
+      (name, param_type, value))
+
+  # Some callers use None, for which we can't do any casting/checking. :(
+  if issubclass(param_type, type(None)):
+    return value
+
+  # Avoid converting a non-string type to a string.
+  if (issubclass(param_type, (six.string_types, six.binary_type)) and
+      not isinstance(value, (six.string_types, six.binary_type))):
+    raise ValueError(fail_msg)
+
+  # Avoid converting a number or string type to a boolean or vice versa.
+  if issubclass(param_type, bool) != isinstance(value, bool):
+    raise ValueError(fail_msg)
+
+  # Avoid converting float to an integer (the reverse is fine).
+  if (issubclass(param_type, numbers.Integral) and
+      not isinstance(value, numbers.Integral)):
+    raise ValueError(fail_msg)
+
+  # Avoid converting a non-numeric type to a numeric type.
+  if (issubclass(param_type, numbers.Number) and
+      not isinstance(value, numbers.Number)):
+    raise ValueError(fail_msg)
+
+  return param_type(value)
+
+
+def parse_values(values, type_map):
+  """Parses hyperparameter values from a string into a python map.
+
+  `values` is a string containing comma-separated `name=value` pairs.
+  For each pair, the value of the hyperparameter named `name` is set to
+  `value`.
+
+  If a hyperparameter name appears multiple times in `values`, a ValueError
+  is raised (e.g. 'a=1,a=2', 'a[1]=1,a[1]=2').
+
+  If a hyperparameter name in both an index assignment and scalar assignment,
+  a ValueError is raised.  (e.g. 'a=[1,2,3],a[0] = 1').
+
+  The hyperparameter name may contain '.' symbols, which will result in an
+  attribute name that is only accessible through the getattr and setattr
+  functions.  (And must be first explicit added through add_hparam.)
+
+  WARNING: Use of '.' in your variable names is allowed, but is not well
+  supported and not recommended.
+
+  The `value` in `name=value` must follows the syntax according to the
+  type of the parameter:
+
+  *  Scalar integer: A Python-parsable integer point value.  E.g.: 1,
+     100, -12.
+  *  Scalar float: A Python-parsable floating point value.  E.g.: 1.0,
+     -.54e89.
+  *  Boolean: Either true or false.
+  *  Scalar string: A non-empty sequence of characters, excluding comma,
+     spaces, and square brackets.  E.g.: foo, bar_1.
+  *  List: A comma separated list of scalar values of the parameter type
+     enclosed in square brackets.  E.g.: [1,2,3], [1.0,1e-12], [high,low].
+
+  When index assignment is used, the corresponding type_map key should be the
+  list name.  E.g. for "arr[1]=0" the type_map must have the key "arr" (not
+  "arr[1]").
+
+  Args:
+    values: String.  Comma separated list of `name=value` pairs where
+      'value' must follow the syntax described above.
+    type_map: A dictionary mapping hyperparameter names to types.  Note every
+      parameter name in values must be a key in type_map.  The values must
+      conform to the types indicated, where a value V is said to conform to a
+      type T if either V has type T, or V is a list of elements of type T.
+      Hence, for a multidimensional parameter 'x' taking float values,
+      'x=[0.1,0.2]' will parse successfully if type_map['x'] = float.
+
+  Returns:
+    A python map mapping each name to either:
+    * A scalar value.
+    * A list of scalar values.
+    * A dictionary mapping index numbers to scalar values.
+    (e.g. "x=5,L=[1,2],arr[1]=3" results in {'x':5,'L':[1,2],'arr':{1:3}}")
+
+  Raises:
+    ValueError: If there is a problem with input.
+    * If `values` cannot be parsed.
+    * If a list is assigned to a list index (e.g. 'a[1] = [1,2,3]').
+    * If the same rvalue is assigned two different values (e.g. 'a=1,a=2',
+      'a[1]=1,a[1]=2', or 'a=1,a=[1]')
+  """
+  results_dictionary = {}
+  pos = 0
+  while pos < len(values):
+    m = PARAM_RE.match(values, pos)
+    if not m:
+      raise ValueError('Malformed hyperparameter value: %s' % values[pos:])
+    # Check that there is a comma between parameters and move past it.
+    pos = m.end()
+    # Parse the values.
+    m_dict = m.groupdict()
+    name = m_dict['name']
+    if name not in type_map:
+      raise ValueError('Unknown hyperparameter type for %s' % name)
+    type_ = type_map[name]
+
+    # Set up correct parsing function (depending on whether type_ is a bool)
+    if type_ == bool:
+
+      def parse_bool(value):
+        if value in ['true', 'True']:
+          return True
+        elif value in ['false', 'False']:
+          return False
+        else:
+          try:
+            return bool(int(value))
+          except ValueError:
+            _parse_fail(name, type_, value, values)
+
+      parse = parse_bool
+    else:
+      parse = type_
+
+    # If a singe value is provided
+    if m_dict['val'] is not None:
+      _process_scalar_value(name, parse, type_, m_dict, values,
+                            results_dictionary)
+
+    # If the assigned value is a list:
+    elif m_dict['vals'] is not None:
+      _process_list_value(name, parse, type_, m_dict, values,
+                          results_dictionary)
+
+    else:  # Not assigned a list or value
+      _parse_fail(name, type_, '', values)
+
+  return results_dictionary
+
+
+class HParams(object):
+  """Class to hold a set of hyperparameters as name-value pairs.
+
+  A `HParams` object holds hyperparameters used to build and train a model,
+  such as the number of hidden units in a neural net layer or the learning rate
+  to use when training.
+
+  You first create a `HParams` object by specifying the names and values of the
+  hyperparameters.
+
+  To make them easily accessible the parameter names are added as direct
+  attributes of the class.  A typical usage is as follows:
+
+  ```python
+  # Create a HParams object specifying names and values of the model
+  # hyperparameters:
+  hparams = HParams(learning_rate=0.1, num_hidden_units=100)
+
+  # The hyperparameter are available as attributes of the HParams object:
+  hparams.learning_rate ==> 0.1
+  hparams.num_hidden_units ==> 100
+  ```
+
+  Hyperparameters have type, which is inferred from the type of their value
+  passed at construction type.   The currently supported types are: integer,
+  float, boolean, string, and list of integer, float, boolean, or string.
+
+  You can override hyperparameter values by calling the
+  [`parse()`](#HParams.parse) method, passing a string of comma separated
+  `name=value` pairs.  This is intended to make it possible to override
+  any hyperparameter values from a single command-line flag to which
+  the user passes 'hyper-param=value' pairs.  It avoids having to define
+  one flag for each hyperparameter.
+
+  The syntax expected for each value depends on the type of the parameter.
+  See `parse()` for a description of the syntax.
+
+  Example:
+
+  ```python
+  # Define a command line flag to pass name=value pairs.
+  # For example using argparse:
+  import argparse
+  parser = argparse.ArgumentParser(description='Train my model.')
+  parser.add_argument('--hparams', type=str,
+                      help='Comma separated list of "name=value" pairs.')
+  args = parser.parse_args()
+  ...
+  def my_program():
+    # Create a HParams object specifying the names and values of the
+    # model hyperparameters:
+    hparams = tf.HParams(learning_rate=0.1, num_hidden_units=100,
+                         activations=['relu', 'tanh'])
+
+    # Override hyperparameters values by parsing the command line
+    hparams.parse(args.hparams)
+
+    # If the user passed `--hparams=learning_rate=0.3` on the command line
+    # then 'hparams' has the following attributes:
+    hparams.learning_rate ==> 0.3
+    hparams.num_hidden_units ==> 100
+    hparams.activations ==> ['relu', 'tanh']
+
+    # If the hyperparameters are in json format use parse_json:
+    hparams.parse_json('{"learning_rate": 0.3, "activations": "relu"}')
+  ```
+  """
+
+  _HAS_DYNAMIC_ATTRIBUTES = True  # Required for pytype checks.
+
+  def __init__(self, hparam_def=None, model_structure=None, **kwargs):
+    """Create an instance of `HParams` from keyword arguments.
+
+    The keyword arguments specify name-values pairs for the hyperparameters.
+    The parameter types are inferred from the type of the values passed.
+
+    The parameter names are added as attributes of `HParams` object, so they
+    can be accessed directly with the dot notation `hparams._name_`.
+
+    Example:
+
+    ```python
+    # Define 3 hyperparameters: 'learning_rate' is a float parameter,
+    # 'num_hidden_units' an integer parameter, and 'activation' a string
+    # parameter.
+    hparams = tf.HParams(
+        learning_rate=0.1, num_hidden_units=100, activation='relu')
+
+    hparams.activation ==> 'relu'
+    ```
+
+    Note that a few names are reserved and cannot be used as hyperparameter
+    names.  If you use one of the reserved name the constructor raises a
+    `ValueError`.
+
+    Args:
+      hparam_def: Serialized hyperparameters, encoded as a hparam_pb2.HParamDef
+        protocol buffer. If provided, this object is initialized by
+        deserializing hparam_def.  Otherwise **kwargs is used.
+      model_structure: An instance of ModelStructure, defining the feature
+        crosses to be used in the Trial.
+      **kwargs: Key-value pairs where the key is the hyperparameter name and
+        the value is the value for the parameter.
+
+    Raises:
+      ValueError: If both `hparam_def` and initialization values are provided,
+        or if one of the arguments is invalid.
+
+    """
+    # Register the hyperparameters and their type in _hparam_types.
+    # This simplifies the implementation of parse().
+    # _hparam_types maps the parameter name to a tuple (type, bool).
+    # The type value is the type of the parameter for scalar hyperparameters,
+    # or the type of the list elements for multidimensional hyperparameters.
+    # The bool value is True if the value is a list, False otherwise.
+    self._hparam_types = {}
+    self._model_structure = model_structure
+    if hparam_def:
+##       self._init_from_proto(hparam_def)
+##       if kwargs:
+##         raise ValueError('hparam_def and initialization values are '
+##                          'mutually exclusive')
+      raise ValueError('hparam_def has been disabled in this version')
+    else:
+      for name, value in six.iteritems(kwargs):
+        self.add_hparam(name, value)
+
+##   def _init_from_proto(self, hparam_def):
+##     """Creates a new HParams from `HParamDef` protocol buffer.
+## 
+##     Args:
+##       hparam_def: `HParamDef` protocol buffer.
+##     """
+##     assert isinstance(hparam_def, hparam_pb2.HParamDef)
+##     for name, value in hparam_def.hparam.items():
+##       kind = value.WhichOneof('kind')
+##       if kind.endswith('_value'):
+##         # Single value.
+##         if kind.startswith('int64'):
+##           # Setting attribute value to be 'int' to ensure the type is compatible
+##           # with both Python2 and Python3.
+##           self.add_hparam(name, int(getattr(value, kind)))
+##         elif kind.startswith('bytes'):
+##           # Setting attribute value to be 'str' to ensure the type is compatible
+##           # with both Python2 and Python3. UTF-8 encoding is assumed.
+##           self.add_hparam(name, compat.as_str(getattr(value, kind)))
+##         else:
+##           self.add_hparam(name, getattr(value, kind))
+##       else:
+##         # List of values.
+##         if kind.startswith('int64'):
+##           # Setting attribute value to be 'int' to ensure the type is compatible
+##           # with both Python2 and Python3.
+##           self.add_hparam(name, [int(v) for v in getattr(value, kind).value])
+##         elif kind.startswith('bytes'):
+##           # Setting attribute value to be 'str' to ensure the type is compatible
+##           # with both Python2 and Python3. UTF-8 encoding is assumed.
+##           self.add_hparam(
+##               name, [compat.as_str(v) for v in getattr(value, kind).value])
+##         else:
+##           self.add_hparam(name, [v for v in getattr(value, kind).value])
+
+  def add_hparam(self, name, value):
+    """Adds {name, value} pair to hyperparameters.
+
+    Args:
+      name: Name of the hyperparameter.
+      value: Value of the hyperparameter. Can be one of the following types:
+        int, float, string, int list, float list, or string list.
+
+    Raises:
+      ValueError: if one of the arguments is invalid.
+    """
+    # Keys in kwargs are unique, but 'name' could the name of a pre-existing
+    # attribute of this object.  In that case we refuse to use it as a
+    # hyperparameter name.
+    if getattr(self, name, None) is not None:
+      raise ValueError('Hyperparameter name is reserved: %s' % name)
+    if isinstance(value, (list, tuple)):
+      if not value:
+        raise ValueError(
+            'Multi-valued hyperparameters cannot be empty: %s' % name)
+      self._hparam_types[name] = (type(value[0]), True)
+    else:
+      self._hparam_types[name] = (type(value), False)
+    setattr(self, name, value)
+
+  def set_hparam(self, name, value):
+    """Set the value of an existing hyperparameter.
+
+    This function verifies that the type of the value matches the type of the
+    existing hyperparameter.
+
+    Args:
+      name: Name of the hyperparameter.
+      value: New value of the hyperparameter.
+
+    Raises:
+      ValueError: If there is a type mismatch.
+    """
+    param_type, is_list = self._hparam_types[name]
+    if isinstance(value, list):
+      if not is_list:
+        raise ValueError(
+            'Must not pass a list for single-valued parameter: %s' % name)
+      setattr(self, name, [
+          _cast_to_type_if_compatible(name, param_type, v) for v in value])
+    else:
+      if is_list:
+        raise ValueError(
+            'Must pass a list for multi-valued parameter: %s.' % name)
+      setattr(self, name, _cast_to_type_if_compatible(name, param_type, value))
+
+  def del_hparam(self, name):
+    """Removes the hyperparameter with key 'name'.
+
+    Args:
+      name: Name of the hyperparameter.
+    """
+    if hasattr(self, name):
+      delattr(self, name)
+      del self._hparam_types[name]
+
+  def parse(self, values):
+    """Override hyperparameter values, parsing new values from a string.
+
+    See parse_values for more detail on the allowed format for values.
+
+    Args:
+      values: String.  Comma separated list of `name=value` pairs where
+        'value' must follow the syntax described above.
+
+    Returns:
+      The `HParams` instance.
+
+    Raises:
+      ValueError: If `values` cannot be parsed.
+    """
+    type_map = dict()
+    for name, t in self._hparam_types.items():
+      param_type, _ = t
+      type_map[name] = param_type
+
+    values_map = parse_values(values, type_map)
+    return self.override_from_dict(values_map)
+
+  def override_from_dict(self, values_dict):
+    """Override hyperparameter values, parsing new values from a dictionary.
+
+    Args:
+      values_dict: Dictionary of name:value pairs.
+
+    Returns:
+      The `HParams` instance.
+
+    Raises:
+      ValueError: If `values_dict` cannot be parsed.
+    """
+    for name, value in values_dict.items():
+      self.set_hparam(name, value)
+    return self
+
+##   @deprecation.deprecated(None, 'Use `override_from_dict`.')
+  def set_from_map(self, values_map):
+    """DEPRECATED. Use override_from_dict."""
+    return self.override_from_dict(values_dict=values_map)
+
+  def set_model_structure(self, model_structure):
+    self._model_structure = model_structure
+
+  def get_model_structure(self):
+    return self._model_structure
+
+  def to_json(self, indent=None, separators=None, sort_keys=False):
+    """Serializes the hyperparameters into JSON.
+
+    Args:
+      indent: If a non-negative integer, JSON array elements and object members
+        will be pretty-printed with that indent level. An indent level of 0, or
+        negative, will only insert newlines. `None` (the default) selects the
+        most compact representation.
+      separators: Optional `(item_separator, key_separator)` tuple. Default is
+        `(', ', ': ')`.
+      sort_keys: If `True`, the output dictionaries will be sorted by key.
+
+    Returns:
+      A JSON string.
+    """
+    return json.dumps(
+        self.values(),
+        indent=indent,
+        separators=separators,
+        sort_keys=sort_keys)
+
+  def parse_json(self, values_json):
+    """Override hyperparameter values, parsing new values from a json object.
+
+    Args:
+      values_json: String containing a json object of name:value pairs.
+
+    Returns:
+      The `HParams` instance.
+
+    Raises:
+      ValueError: If `values_json` cannot be parsed.
+    """
+    values_map = json.loads(values_json)
+    return self.override_from_dict(values_map)
+
+  def values(self):
+    """Return the hyperparameter values as a Python dictionary.
+
+    Returns:
+      A dictionary with hyperparameter names as keys.  The values are the
+      hyperparameter values.
+    """
+    return {n: getattr(self, n) for n in self._hparam_types.keys()}
+
+  def get(self, key, default=None):
+    """Returns the value of `key` if it exists, else `default`."""
+    if key in self._hparam_types:
+      # Ensure that default is compatible with the parameter type.
+      if default is not None:
+        param_type, is_param_list = self._hparam_types[key]
+        type_str = 'list<%s>' % param_type if is_param_list else str(param_type)
+        fail_msg = ("Hparam '%s' of type '%s' is incompatible with "
+                    'default=%s' % (key, type_str, default))
+
+        is_default_list = isinstance(default, list)
+        if is_param_list != is_default_list:
+          raise ValueError(fail_msg)
+
+        try:
+          if is_default_list:
+            for value in default:
+              _cast_to_type_if_compatible(key, param_type, value)
+          else:
+            _cast_to_type_if_compatible(key, param_type, default)
+        except ValueError as e:
+          raise ValueError('%s. %s' % (fail_msg, e))
+
+      return getattr(self, key)
+
+    return default
+
+  def __contains__(self, key):
+    return key in self._hparam_types
+
+  def __str__(self):
+    return str(sorted(self.values().items()))
+
+  def __repr__(self):
+    return '%s(%s)' % (type(self).__name__, self.__str__())
+
+  @staticmethod
+  def _get_kind_name(param_type, is_list):
+    """Returns the field name given parameter type and is_list.
+
+    Args:
+      param_type: Data type of the hparam.
+      is_list: Whether this is a list.
+
+    Returns:
+      A string representation of the field name.
+
+    Raises:
+      ValueError: If parameter type is not recognized.
+    """
+    if issubclass(param_type, bool):
+      # This check must happen before issubclass(param_type, six.integer_types),
+      # since Python considers bool to be a subclass of int.
+      typename = 'bool'
+    elif issubclass(param_type, six.integer_types):
+      # Setting 'int' and 'long' types to be 'int64' to ensure the type is
+      # compatible with both Python2 and Python3.
+      typename = 'int64'
+    elif issubclass(param_type, (six.string_types, six.binary_type)):
+      # Setting 'string' and 'bytes' types to be 'bytes' to ensure the type is
+      # compatible with both Python2 and Python3.
+      typename = 'bytes'
+    elif issubclass(param_type, float):
+      typename = 'float'
+    else:
+      raise ValueError('Unsupported parameter type: %s' % str(param_type))
+
+    suffix = 'list' if is_list else 'value'
+    return '_'.join([typename, suffix])
+
+##   def to_proto(self, export_scope=None):  # pylint: disable=unused-argument
+##     """Converts a `HParams` object to a `HParamDef` protocol buffer.
+## 
+##     Args:
+##       export_scope: Optional `string`. Name scope to remove.
+## 
+##     Returns:
+##       A `HParamDef` protocol buffer.
+##     """
+##     hparam_proto = hparam_pb2.HParamDef()
+##     for name in self._hparam_types:
+##       # Parse the values.
+##       param_type, is_list = self._hparam_types.get(name, (None, None))
+##       kind = HParams._get_kind_name(param_type, is_list)
+## 
+##       if is_list:
+##         if kind.startswith('bytes'):
+##           v_list = [compat.as_bytes(v) for v in getattr(self, name)]
+##         else:
+##           v_list = [v for v in getattr(self, name)]
+##         getattr(hparam_proto.hparam[name], kind).value.extend(v_list)
+##       else:
+##         v = getattr(self, name)
+##         if kind.startswith('bytes'):
+##           v = compat.as_bytes(getattr(self, name))
+##         setattr(hparam_proto.hparam[name], kind, v)
+## 
+##     return hparam_proto
+
+##   @staticmethod
+##   def from_proto(hparam_def, import_scope=None):  # pylint: disable=unused-argument
+##     return HParams(hparam_def=hparam_def)
+
+
+## ops.register_proto_function(
+##     'hparams',
+##     proto_type=hparam_pb2.HParamDef,
+##     to_proto=HParams.to_proto,
+##     from_proto=HParams.from_proto)

tfcompat/readme.md

@@ -0,0 +1,8 @@
+Source: hparam.py copied from tensorflow v1.12.0.
+
+https://github.com/tensorflow/tensorflow/blob/v1.12.0/tensorflow/contrib/training/python/training/hparam.py
+
+with the following:
+wget https://github.com/tensorflow/tensorflow/raw/v1.12.0/tensorflow/contrib/training/python/training/hparam.py
+
+Once all other tensorflow dependencies of these file are removed, the class keeps its goal. Functions not available due to this process are not used in this project.

utils.py

@@ -0,0 +1,111 @@
+import copy
+import torch
+import numpy as np
+from scipy import signal
+from librosa.filters import mel
+from scipy.signal import get_window
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def butter_highpass(cutoff, fs, order=5):
+    nyq = 0.5 * fs
+    normal_cutoff = cutoff / nyq
+    b, a = signal.butter(order, normal_cutoff, btype='high', analog=False)
+    return b, a
+    
+    
+    
+def pySTFT(x, fft_length=1024, hop_length=256):
+    
+    x = np.pad(x, int(fft_length//2), mode='reflect')
+    
+    noverlap = fft_length - hop_length
+    shape = x.shape[:-1]+((x.shape[-1]-noverlap)//hop_length, fft_length)
+    strides = x.strides[:-1]+(hop_length*x.strides[-1], x.strides[-1])
+    result = np.lib.stride_tricks.as_strided(x, shape=shape,
+                                             strides=strides)
+    
+    fft_window = get_window('hann', fft_length, fftbins=True)
+    result = np.fft.rfft(fft_window * result, n=fft_length).T
+    
+    return np.abs(result)
+
+
+
+class LinearNorm(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
+        super(LinearNorm, self).__init__()
+        self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
+
+        torch.nn.init.xavier_uniform_(
+            self.linear_layer.weight,
+            gain=torch.nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+
+class ConvNorm(torch.nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1,
+                 padding=None, dilation=1, bias=True, w_init_gain='linear'):
+        super(ConvNorm, self).__init__()
+        if padding is None:
+            assert(kernel_size % 2 == 1)
+            padding = int(dilation * (kernel_size - 1) / 2)
+
+        self.conv = torch.nn.Conv1d(in_channels, out_channels,
+                                    kernel_size=kernel_size, stride=stride,
+                                    padding=padding, dilation=dilation,
+                                    bias=bias)
+
+        torch.nn.init.xavier_uniform_(
+            self.conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, signal):
+        conv_signal = self.conv(signal)
+        return conv_signal
+    
+    
+
+def filter_bank_mean(num_rep, codes_mask, max_len_long):
+    '''
+    num_rep (B, L)
+    codes_mask (B, L)
+    
+    output: filterbank (B, L, max_len_fake)
+    
+    zero pad in codes must be real zero
+    '''
+    
+    num_rep = num_rep.unsqueeze(-1) # (B, L, 1)
+    codes_mask = codes_mask.unsqueeze(-1) # (B, L, 1)
+    num_rep = num_rep * codes_mask
+    
+    right_edge = num_rep.cumsum(dim=1)
+    left_edge = torch.zeros_like(right_edge)
+    left_edge[:, 1:, :] = right_edge[:, :-1, :]
+    right_edge = right_edge.ceil()
+    left_edge = left_edge.floor()
+    
+    index = torch.arange(1, max_len_long+1, device=num_rep.device).view(1, 1, -1)
+    
+    lower = index - left_edge
+
+    right_edge_flip = max_len_long - right_edge
+    
+    upper = (index - right_edge_flip).flip(dims=(2,))
+    
+    # triangular pooling
+    fb = F.relu(torch.min(lower, upper)).float()
+    
+    # mean pooling
+    fb = (fb > 0).float()
+    
+    norm = fb.sum(dim=-1, keepdim=True)
+    norm[norm==0] = 1.0
+    
+    fb = fb / norm
+    
+    return fb * codes_mask