update model & inference

- replace amsgrad with sgd+momentum (beta=0.9)
- set lr=0.2, and decay it by half every 32 epochs
- shrink h from 256 to 192
- separate hidden and latent, and set h_latent=64
- apply gradient clipping by L2 norm (max_norm=1)
- increase epochs from 160 to 192
- add tqdm progress bar display
- integrate incremental generation during inference

decoder.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3887ea54b29b8d19669b4020dff46804593262ec49ef159816925274fc418572
-size 3925384
+oid sha256:deeea664de143a71c87e67ba2af78aa88320fcd401c2c12a40183060f78b0e15
+size 2078336

inference.py

@@ -6,16 +6,21 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 class DecoderGRU(nn.Module):
-    def __init__(self, hidden_size, output_size):
+    def __init__(self, latent_size, hidden_size, output_size):
         super(DecoderGRU, self).__init__()
-        self.proj = nn.Linear(hidden_size, 2 * hidden_size)
+        self.proj1 = nn.Linear(latent_size, latent_size)
+        self.proj_activation = nn.ReLU()
+        self.proj2 = nn.Linear(latent_size, 2 * hidden_size)
         self.embedding = nn.Embedding(output_size, hidden_size)
         self.gru = nn.GRU(hidden_size, hidden_size, num_layers=2, batch_first=True)
         self.out = nn.Linear(hidden_size, output_size)
 
     def forward(self, encoder_sample, target_tensor=None, max_length=16):
         batch_size = encoder_sample.size(0)
-        decoder_hidden = self.proj(encoder_sample).view(batch_size, 2, -1).permute(1, 0, 2).contiguous()
+        decoder_hidden = self.proj1(encoder_sample)
+        decoder_hidden = self.proj_activation(decoder_hidden)
+        decoder_hidden = self.proj2(decoder_hidden)
+        decoder_hidden = decoder_hidden.view(batch_size, 2, -1).permute(1, 0, 2).contiguous()
         if target_tensor is not None:
             decoder_input = target_tensor
             decoder_outputs, decoder_hidden = self.forward_step(decoder_input, decoder_hidden)
@@ -46,8 +51,9 @@ katakana = list('゠ァアィイゥウェエォオカガキギクグケゲコゴ
 vocab = ['<pad>', '<sos>', '<eos>'] + katakana
 vocab_dict = {v: k for k, v in enumerate(vocab)}
 
-h=256
+h_latent=64
 max_len=40
+names=16
 
 def detokenize(tokens):
     if EOS_token in tokens:
@@ -55,6 +61,9 @@ def detokenize(tokens):
     else:
         return None
 
-for name in [detokenize(seq) for seq in dec(torch.randn(16,h), max_length=max_len)[0].topk(1)[1].squeeze().tolist()]:
-    if name is not None:
-        print(name)
+while True:
+    print('generating names...')
+    for name in [detokenize(seq) for seq in dec(torch.randn(names,h_latent), max_length=max_len)[0].topk(1)[1].squeeze().tolist()]:
+        if name is not None:
+            print(name)
+    input("press enter to continue generation...")

model.py

@@ -5,6 +5,8 @@ import torch
 import torch.nn as nn
 from torch import optim
 from torch.utils.data import DataLoader, Dataset
+from torch.optim.lr_scheduler import StepLR
+from tqdm.auto import tqdm
 import torch.nn.functional as F
 import pandas as pd
 
@@ -19,11 +21,16 @@ vocab_dict = {v: k for k, v in enumerate(vocab)}
 
 texts = pd.read_csv('rolename.txt', header=None)[0].tolist()
 vocab_size=len(vocab)
-h=256
+h=192
+h_latent=64
 max_len=40
 bs=128
-lr=5e-4
-epochs=160
+lr=0.2
+lr_step_size=32
+lr_decay=0.5
+momentum=0.9
+epochs=192
+grad_max_norm=1
 
 def tokenize(text):
     return [vocab_dict[ch] for ch in text]
@@ -50,15 +57,15 @@ class BatchNormVAE(nn.Module): # https://spaces.ac.cn/archives/7381/
         return mu*scale_mu, sigma*scale_sigma
 
 class EncoderVAEBiGRU(nn.Module):
-    def __init__(self, input_size, hidden_size, dropout_p=0.1):
+    def __init__(self, input_size, hidden_size, latent_size, dropout_p=0.1):
         super(EncoderVAEBiGRU, self).__init__()
         self.hidden_size = hidden_size
         self.embedding = nn.Embedding(input_size, hidden_size)
         self.gru = nn.GRU(hidden_size, hidden_size, num_layers=2, batch_first=True, bidirectional=True)
-        self.proj_mu = nn.Linear(4 * hidden_size, hidden_size)
-        self.proj_sigma = nn.Linear(4 * hidden_size, hidden_size)
+        self.proj_mu = nn.Linear(4 * hidden_size, latent_size)
+        self.proj_sigma = nn.Linear(4 * hidden_size, latent_size)
         self.dropout = nn.Dropout(dropout_p)
-        self.bn = BatchNormVAE(hidden_size)
+        self.bn = BatchNormVAE(latent_size)
 
     def forward(self, input, input_lengths):
         input_lengths = input_lengths.to('cpu')
@@ -76,16 +83,21 @@ class EncoderVAEBiGRU(nn.Module):
         return eps * sigma + mu # var is sigma^2
 
 class DecoderGRU(nn.Module):
-    def __init__(self, hidden_size, output_size):
+    def __init__(self, latent_size, hidden_size, output_size):
         super(DecoderGRU, self).__init__()
-        self.proj = nn.Linear(hidden_size, 2 * hidden_size)
+        self.proj1 = nn.Linear(latent_size, latent_size)
+        self.proj_activation = nn.ReLU()
+        self.proj2 = nn.Linear(latent_size, 2 * hidden_size)
         self.embedding = nn.Embedding(output_size, hidden_size)
         self.gru = nn.GRU(hidden_size, hidden_size, num_layers=2, batch_first=True)
         self.out = nn.Linear(hidden_size, output_size)
 
     def forward(self, encoder_sample, target_tensor=None, max_length=16):
         batch_size = encoder_sample.size(0)
-        decoder_hidden = self.proj(encoder_sample).view(batch_size, 2, -1).permute(1, 0, 2).contiguous()
+        decoder_hidden = self.proj1(encoder_sample)
+        decoder_hidden = self.proj_activation(decoder_hidden)
+        decoder_hidden = self.proj2(decoder_hidden)
+        decoder_hidden = decoder_hidden.view(batch_size, 2, -1).permute(1, 0, 2).contiguous()
         if target_tensor is not None:
             decoder_input = target_tensor
             decoder_outputs, decoder_hidden = self.forward_step(decoder_input, decoder_hidden)
@@ -136,7 +148,7 @@ dataloader = DataLoader(
     generator=torch.Generator(device='cuda'),
 )
 
-def train_epoch(dataloader, encoder, decoder, optimizer):
+def train_epoch(dataloader, encoder, decoder, optimizer, max_norm, norm_p=2):
     total_loss = 0
     nll = nn.NLLLoss()
     for enc_text, enc_len, input_text, target_text in dataloader:
@@ -150,19 +162,25 @@ def train_epoch(dataloader, encoder, decoder, optimizer):
         loss = loss_recons + loss_kld
         loss.backward()
 
+        # gradient clipping by norm
+        nn.utils.clip_grad_norm_(list(encoder.parameters()) + list(decoder.parameters()), max_norm, norm_type=norm_p)
+
         optimizer.step()
 
         total_loss += loss.item()
     return total_loss / len(dataloader)
 
-enc = EncoderVAEBiGRU(vocab_size, h).train()
-dec = DecoderGRU(h, vocab_size).train()
-optimizer = optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=lr, amsgrad=True) # AMSGrad
+encoder = EncoderVAEBiGRU(vocab_size, h, h_latent).train()
+decoder = DecoderGRU(h_latent, h, vocab_size).train()
+optimizer = optim.SGD(list(encoder.parameters()) + list(decoder.parameters()), lr=lr, momentum=momentum) # momentum
+scheduler = StepLR(optimizer, step_size=lr_step_size, gamma=lr_decay)
 
-for i in range(epochs):
-    print('epoch=%d, loss=%f' % (i, train_epoch(dataloader, enc, dec, optimizer)))
+with tqdm(range(epochs), desc='Training') as pbar:
+    for i in pbar:
+        pbar.set_postfix(loss=train_epoch(dataloader, encoder, decoder, optimizer, grad_max_norm))
+        scheduler.step()
 
-dec.eval()
-for name in [detokenize(seq) for seq in dec(torch.randn(8,h), max_length=max_len)[0].topk(1)[1].squeeze().tolist()]:
+decoder.eval()
+for name in [detokenize(seq) for seq in decoder(torch.randn(8,h_latent), max_length=max_len)[0].topk(1)[1].squeeze().tolist()]:
     print(name)
-torch.save(dec, 'decoder.pt')
+torch.save(decoder, 'decoder.pt')