├── assets ├── d_loss.png ├── g_loss.png └── w_distance.png ├── LICENSE ├── README.md ├── utils.py ├── .gitignore ├── model.py └── train.py /assets/d_loss.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/keon/text-wgan/HEAD/assets/d_loss.png -------------------------------------------------------------------------------- /assets/g_loss.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/keon/text-wgan/HEAD/assets/g_loss.png -------------------------------------------------------------------------------- /assets/w_distance.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/keon/text-wgan/HEAD/assets/w_distance.png -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Keon 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # text-wgan 2 | Improved Training of Wasserstein GANs for Text 3 | 4 | WIP! 5 | 6 | Based on the paper [Improved Training of Wasserstein GANs](https://arxiv.org/abs/1704.00028). 7 | 8 | 9 | ![](./assets/g_loss.png) 10 | ![](./assets/d_loss.png) 11 | ![](./assets/w_distance.png) 12 | 13 | Sample after 15000 iterations 14 | 15 | ``` 16 | this shere pmongerin liked braning sith ang and and bod whosd by ic-t andoningl, ames like this to kookt marmangenedcoy wing the lemest celic iating telike harhart to be eyentinlto ok wigical eaitly is lmdoincti isclederling tleads of the vidtlin-anew,.he mostedion has cebiche i 17 | 18 | oover stortly whiting eventa fot and and purvingshing dovarnyy marming, and dood the e.rit ond grases in ane the dannierdey. tris cootied arneges. wrath bad and dcting. made...t suming but dilting actins woat get ooemouthgind good cran tion and d vared intaliit,stming i sad ge ge 19 | 20 | somd. ao, me dencan,i huck on the patfar its doscate must rusel, a deul the pont of hesvel, any me sage movie mame it was allmovce his was poring, evellever vowe risuide a r wancidery hen onemor ie, in that wastany no e sate funmout this moveestink, medes film, is.n, i an pre is 21 | ``` 22 | 23 | ## References 24 | 25 | Based on the following implementations 26 | 27 | - [caogang/wgan-gp](https://github.com/caogang/wgan-gp) 28 | - [igul222/improved_wgan_training](https://github.com/igul222/improved_wgan_training) 29 | - [martinarjovsky/WassersteinGAN](https://github.com/martinarjovsky/WassersteinGAN) 30 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | from torch.autograd import Variable 4 | 5 | 6 | def to_onehot(index, vocab_size): 7 | batch_size, seq_len = index.size(0), index.size(1) 8 | onehot = torch.FloatTensor(batch_size, seq_len, vocab_size).zero_() 9 | onehot.scatter_(2, index.data.cpu().unsqueeze(2), 1) 10 | return onehot 11 | 12 | 13 | def sample(G, TEXT, batch_size, seq_len, vocab_size, use_cuda=True): 14 | noise = torch.randn(batch_size, 128) 15 | if use_cuda: 16 | noise = noise.cuda() 17 | noisev = Variable(noise, volatile=True) 18 | samples = G(noisev) 19 | samples = samples.view(-1, seq_len, vocab_size) 20 | _, argmax = torch.max(samples, 2) 21 | argmax = argmax.cpu().data 22 | decoded_samples = [] 23 | for i in range(len(argmax)): 24 | decoded = "".join([TEXT.vocab.itos[s] for s in argmax[i]]) 25 | decoded_samples.append(decoded) 26 | return decoded_samples 27 | 28 | 29 | def plot(title, vis, x, y, win=None): 30 | if win is None: 31 | win = vis.line( 32 | X=np.asarray([x]), 33 | Y=np.asarray([y]), 34 | opts=dict(title=title, xlabel='Batch', ylabel='Loss') 35 | ) 36 | else: 37 | vis.line(X=np.asarray([x]), Y=np.asarray([y]), 38 | win=win, update='append') 39 | return win 40 | 41 | 42 | def log_sample(title, vis, text, win=None): 43 | if win is None: 44 | win = vis.text(text, opts=dict(title=title)) 45 | else: 46 | vis.text(text, win=win, append=True) 47 | return win 48 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | 49 | # Translations 50 | *.mo 51 | *.pot 52 | 53 | # Django stuff: 54 | *.log 55 | local_settings.py 56 | 57 | # Flask stuff: 58 | instance/ 59 | .webassets-cache 60 | 61 | # Scrapy stuff: 62 | .scrapy 63 | 64 | # Sphinx documentation 65 | docs/_build/ 66 | 67 | # PyBuilder 68 | target/ 69 | 70 | # Jupyter Notebook 71 | .ipynb_checkpoints 72 | 73 | # pyenv 74 | .python-version 75 | 76 | # celery beat schedule file 77 | celerybeat-schedule 78 | 79 | # SageMath parsed files 80 | *.sage.py 81 | 82 | # dotenv 83 | .env 84 | 85 | # virtualenv 86 | .venv 87 | venv/ 88 | ENV/ 89 | 90 | # Spyder project settings 91 | .spyderproject 92 | .spyproject 93 | 94 | # Rope project settings 95 | .ropeproject 96 | 97 | # mkdocs documentation 98 | /site 99 | 100 | # mypy 101 | .mypy_cache/ 102 | 103 | .data 104 | .save 105 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | 3 | 4 | class ResBlock(nn.Module): 5 | def __init__(self, dim): 6 | super(ResBlock, self).__init__() 7 | self.res_block = nn.Sequential( 8 | nn.ReLU(True), 9 | nn.Conv1d(dim, dim, 5, padding=2), # nn.Linear(DIM, DIM), 10 | nn.ReLU(True), 11 | nn.Conv1d(dim, dim, 5, padding=2), # nn.Linear(DIM, DIM), 12 | ) 13 | 14 | def forward(self, input): 15 | output = self.res_block(input) 16 | return input + (0.3 * output) 17 | 18 | 19 | class Generator(nn.Module): 20 | def __init__(self, dim, seq_len, vocab_size): 21 | super(Generator, self).__init__() 22 | self.dim = dim 23 | self.seq_len = seq_len 24 | 25 | self.fc1 = nn.Linear(128, dim * seq_len) 26 | self.block = nn.Sequential( 27 | ResBlock(dim), 28 | ResBlock(dim), 29 | ResBlock(dim), 30 | ResBlock(dim), 31 | ResBlock(dim), 32 | ) 33 | self.conv1 = nn.Conv1d(dim, vocab_size, 1) 34 | self.softmax = nn.Softmax() 35 | 36 | def forward(self, noise): 37 | batch_size = noise.size(0) 38 | output = self.fc1(noise) 39 | # (BATCH_SIZE, DIM, SEQ_LEN) 40 | output = output.view(-1, self.dim, self.seq_len) 41 | output = self.block(output) 42 | output = self.conv1(output) 43 | output = output.transpose(1, 2) 44 | shape = output.size() 45 | output = output.contiguous() 46 | output = output.view(batch_size * self.seq_len, -1) 47 | output = self.softmax(output) 48 | # (BATCH_SIZE, SEQ_LEN, len(charmap)) 49 | return output.view(shape) 50 | 51 | 52 | class Discriminator(nn.Module): 53 | def __init__(self, dim, seq_len, vocab_size): 54 | super(Discriminator, self).__init__() 55 | self.dim = dim 56 | self.seq_len = seq_len 57 | 58 | self.block = nn.Sequential( 59 | ResBlock(dim), 60 | ResBlock(dim), 61 | ResBlock(dim), 62 | ResBlock(dim), 63 | ResBlock(dim), 64 | ) 65 | self.conv1d = nn.Conv1d(vocab_size, dim, 1) 66 | self.linear = nn.Linear(seq_len * dim, 1) 67 | 68 | def forward(self, input): 69 | # (BATCH_SIZE, VOCAB_SIZE, SEQ_LEN) 70 | output = input.transpose(1, 2) 71 | output = self.conv1d(output) 72 | output = self.block(output) 73 | output = output.view(-1, self.seq_len * self.dim) 74 | output = self.linear(output) 75 | return output 76 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import argparse 3 | import torch 4 | from torch.autograd import Variable, grad 5 | from torch import optim 6 | from torchtext.data import Field, BucketIterator 7 | from torchtext.datasets import IMDB 8 | from model import Generator, Discriminator 9 | from visdom import Visdom 10 | from utils import to_onehot, sample, plot, log_sample 11 | 12 | 13 | def parse_arguments(): 14 | p = argparse.ArgumentParser(description='Hyperparams') 15 | p.add_argument('-batchs', type=int, default=500000, 16 | help='number of epochs for train') 17 | p.add_argument('-critic_iters', type=int, default=5, 18 | help='critic iterations') 19 | p.add_argument('-batch_size', type=int, default=8, 20 | help='number of epochs for train') 21 | p.add_argument('-seq_len', type=int, default=280, 22 | help='sequence length') 23 | p.add_argument('-lamb', type=int, default=10, 24 | help='lambda') 25 | p.add_argument('-lr', type=float, default=1e-4, 26 | help='initial learning rate') 27 | return p.parse_args() 28 | 29 | 30 | def penalize_grad(D, real, fake, batch_size, lamb, use_cuda=True): 31 | """ 32 | lamb: lambda 33 | """ 34 | alpha = torch.rand(batch_size, 1, 1).expand(real.size()) 35 | if use_cuda: 36 | alpha = alpha.cuda() 37 | interpolates = alpha * real + ((1 - alpha) * fake) 38 | if use_cuda: 39 | interpolates = interpolates.cuda() 40 | interpolates = Variable(interpolates, requires_grad=True) 41 | d_interpolates = D(interpolates) 42 | ones = torch.ones(d_interpolates.size()) 43 | if use_cuda: 44 | ones = ones.cuda() 45 | gradients = grad(outputs=d_interpolates, inputs=interpolates, 46 | grad_outputs=ones, create_graph=True, 47 | retain_graph=True, only_inputs=True)[0] 48 | grad_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * lamb 49 | return grad_penalty 50 | 51 | 52 | def train_discriminator(D, G, optim_D, real, lamb, batch_size, use_cuda=True): 53 | D.zero_grad() 54 | 55 | # train with real 56 | d_real = D(real) 57 | d_real = d_real.mean() 58 | d_real.backward(mone) 59 | 60 | # train with fake 61 | noise = torch.randn(batch_size, 128) 62 | if use_cuda: 63 | noise = noise.cuda() 64 | noise = Variable(noise, volatile=True) # freeze G 65 | fake = G(noise) 66 | fake = Variable(fake.data) 67 | inputv = fake 68 | d_fake = D(inputv) 69 | d_fake = d_fake.mean() 70 | d_fake.backward(one) 71 | 72 | grad_penalty = penalize_grad(D, real.data, fake.data, 73 | batch_size, lamb, use_cuda) 74 | grad_penalty.backward() 75 | 76 | d_loss = d_fake - d_real + grad_penalty 77 | wasserstein = d_real - d_fake 78 | optim_D.step() 79 | return d_loss, wasserstein 80 | 81 | 82 | def train_generator(D, G, optim_G, batch_size, use_cuda=True): 83 | G.zero_grad() 84 | noise = torch.randn(batch_size, 128) 85 | if use_cuda: 86 | noise = noise.cuda() 87 | noisev = Variable(noise) 88 | fake = G(noisev) 89 | g = D(fake) 90 | g = g.mean() 91 | g.backward(mone) 92 | g_loss = -g 93 | optim_G.step() 94 | return g_loss 95 | 96 | 97 | def main(): 98 | args = parse_arguments() 99 | use_cuda = torch.cuda.is_available() 100 | 101 | # visdom for plotting 102 | vis = Visdom() 103 | win_g, win_d, win_w = None, None, None 104 | assert vis.check_connection() 105 | 106 | # load datasets 107 | print("[!] preparing dataset...") 108 | TEXT = Field(lower=True, fix_length=args.seq_len, 109 | tokenize=list, batch_first=True) 110 | LABEL = Field(sequential=False) 111 | train_data, test_data = IMDB.splits(TEXT, LABEL) 112 | TEXT.build_vocab(train_data) 113 | LABEL.build_vocab(train_data) 114 | train_iter, test_iter = BucketIterator.splits( 115 | (train_data, test_data), batch_size=args.batch_size, repeat=True) 116 | vocab_size = len(TEXT.vocab) 117 | print("[TRAIN]:%d (dataset:%d)\t[TEST]:%d (dataset:%d)\t[VOCAB]:%d" 118 | % (len(train_iter), len(train_iter.dataset), 119 | len(test_iter), len(test_iter.dataset), vocab_size)) 120 | 121 | # instantiate models 122 | G = Generator(dim=512, seq_len=args.seq_len, vocab_size=vocab_size) 123 | D = Discriminator(dim=512, seq_len=args.seq_len, vocab_size=vocab_size) 124 | optim_G = optim.Adam(G.parameters(), lr=args.lr, betas=(0.5, 0.9)) 125 | optim_D = optim.Adam(D.parameters(), lr=args.lr, betas=(0.5, 0.9)) 126 | 127 | global one, mone 128 | one = torch.FloatTensor([1]) 129 | mone = one * -1 130 | if use_cuda: 131 | G, D = G.cuda(), D.cuda() 132 | one, mone = one.cuda(), mone.cuda() 133 | 134 | train_iter = iter(train_iter) 135 | batch_size = args.batch_size 136 | for b in range(1, args.batchs+1): 137 | # (1) Update D network 138 | for p in D.parameters(): # reset requires_grad 139 | p.requires_grad = True 140 | for iter_d in range(args.critic_iters): # CRITIC_ITERS 141 | batch = next(train_iter) 142 | text, label = batch.text, batch.label 143 | text = to_onehot(text, vocab_size) 144 | if use_cuda: 145 | text = text.cuda() 146 | real = Variable(text) 147 | d_loss, wasserstein = train_discriminator( 148 | D, G, optim_D, real, args.lamb, batch_size, use_cuda) 149 | # (2) Update G network 150 | for p in D.parameters(): 151 | p.requires_grad = False # to avoid computation 152 | g_loss = train_generator(D, G, optim_G, batch_size, use_cuda) 153 | 154 | # plot losses on visdom 155 | win_d = plot('Discriminator Loss', vis, 156 | x=b, y=d_loss.data[0], win=win_d) 157 | win_g = plot('Generator Loss', vis, 158 | x=b, y=g_loss.data[0], win=win_g) 159 | win_w = plot('Wasserstein Distance', vis, 160 | x=b, y=wasserstein.data[0], win=win_w) 161 | 162 | if b % 500 == 0 and b > 1: 163 | samples = sample(G, TEXT, 1, args.seq_len, vocab_size, use_cuda) 164 | print("[%d] D:%5.2f G:%5.2f W:%5.2f \nsample:%s \t [%d]" % 165 | (b, d_loss.data[0], g_loss.data[0], wasserstein.data[0], 166 | samples[0], label.data[0])) 167 | log_sample("Sample %d" % b, vis, samples) 168 | if b % 5000 == 0 and b > 1: 169 | print("[!] saving model") 170 | if not os.path.isdir(".save"): 171 | os.makedirs(".save") 172 | torch.save(G.state_dict(), './.save/wgan_g_%d.pt' % (b)) 173 | torch.save(D.state_dict(), './.save/wgan_d_%d.pt' % (b)) 174 | 175 | 176 | if __name__ == "__main__": 177 | try: 178 | main() 179 | except KeyboardInterrupt as e: 180 | print("[STOP]", e) 181 | --------------------------------------------------------------------------------