├── .travis.yml ├── DCGAN ├── DCGAN.py ├── __init__.py └── run_DCGAN.py ├── GAN ├── GAN.py ├── __init__.py └── run_GAN.py ├── ImprovedGAN ├── ImprovedGAN.py ├── __init__.py └── run_ImprovedGAN.py ├── InfoGAN ├── InfoGAN.py ├── __init__.py └── run_InfoGAN.py ├── LAPGAN ├── LAPGAN.py ├── __init__.py └── run_LAPGAN.py ├── README.md ├── images ├── DCGAN.png ├── ImprovedGAN.png ├── InfoGAN.png └── real.png ├── requirements.txt └── tests ├── __init__.py └── test_GANs.py /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | cache: pip 3 | dist: trusty 4 | sudo: required 5 | 6 | 7 | python: 8 | - "2.7" 9 | 10 | 11 | before_install: 12 | - pip install -U pip 13 | - pip install wheel 14 | - pip install coveralls 15 | - sudo apt-get update 16 | 17 | 18 | env: 19 | global: 20 | - PIP_WHEEL_DIR=$HOME/.cache/pip/wheels 21 | - PIP_FIND_LINKS=file://$HOME/.cache/pip/wheels 22 | 23 | 24 | install: 25 | - pip wheel -r requirements.txt 26 | - pip install -r requirements.txt 27 | - pip install https://download.pytorch.org/whl/cu75/torch-0.1.10.post2-cp27-none-linux_x86_64.whl 28 | - pip install torchvision 29 | 30 | 31 | script: 32 | - py.test . --cov=./ 33 | - flake8 ./ 34 | 35 | 36 | after_success: 37 | - coveralls 38 | -------------------------------------------------------------------------------- /DCGAN/DCGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | class DCGAN_Discriminator(nn.Module): 9 | 10 | def __init__(self, featmap_dim=512, n_channel=1): 11 | super(DCGAN_Discriminator, self).__init__() 12 | self.featmap_dim = featmap_dim 13 | self.conv1 = nn.Conv2d(n_channel, featmap_dim / 4, 5, 14 | stride=2, padding=2) 15 | 16 | self.conv2 = nn.Conv2d(featmap_dim / 4, featmap_dim / 2, 5, 17 | stride=2, padding=2) 18 | self.BN2 = nn.BatchNorm2d(featmap_dim / 2) 19 | 20 | self.conv3 = nn.Conv2d(featmap_dim / 2, featmap_dim, 5, 21 | stride=2, padding=2) 22 | self.BN3 = nn.BatchNorm2d(featmap_dim) 23 | 24 | self.fc = nn.Linear(featmap_dim * 4 * 4, 1) 25 | 26 | def forward(self, x): 27 | """ 28 | Strided convulation layers, 29 | Batch Normalization after convulation but not at input layer, 30 | LeakyReLU activation function with slope 0.2. 31 | """ 32 | x = F.leaky_relu(self.conv1(x), negative_slope=0.2) 33 | x = F.leaky_relu(self.BN2(self.conv2(x)), negative_slope=0.2) 34 | x = F.leaky_relu(self.BN3(self.conv3(x)), negative_slope=0.2) 35 | x = x.view(-1, self.featmap_dim * 4 * 4) 36 | x = F.sigmoid(self.fc(x)) 37 | return x 38 | 39 | 40 | class DCGAN_Generator(nn.Module): 41 | 42 | def __init__(self, featmap_dim=1024, n_channel=1, noise_dim=100): 43 | super(DCGAN_Generator, self).__init__() 44 | self.featmap_dim = featmap_dim 45 | self.fc1 = nn.Linear(noise_dim, 4 * 4 * featmap_dim) 46 | self.conv1 = nn.ConvTranspose2d(featmap_dim, (featmap_dim / 2), 5, 47 | stride=2, padding=2) 48 | 49 | self.BN1 = nn.BatchNorm2d(featmap_dim / 2) 50 | self.conv2 = nn.ConvTranspose2d(featmap_dim / 2, featmap_dim / 4, 6, 51 | stride=2, padding=2) 52 | 53 | self.BN2 = nn.BatchNorm2d(featmap_dim / 4) 54 | self.conv3 = nn.ConvTranspose2d(featmap_dim / 4, n_channel, 6, 55 | stride=2, padding=2) 56 | 57 | def forward(self, x): 58 | """ 59 | Project noise to featureMap * width * height, 60 | Batch Normalization after convulation but not at output layer, 61 | ReLU activation function. 62 | """ 63 | x = self.fc1(x) 64 | x = x.view(-1, self.featmap_dim, 4, 4) 65 | x = F.relu(self.BN1(self.conv1(x))) 66 | x = F.relu(self.BN2(self.conv2(x))) 67 | x = F.tanh(self.conv3(x)) 68 | 69 | return x 70 | -------------------------------------------------------------------------------- /DCGAN/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /DCGAN/run_DCGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch 5 | import torch.optim as optim 6 | import torchvision 7 | import torchvision.transforms as transforms 8 | import numpy as np 9 | 10 | from torch.autograd import Variable 11 | from DCGAN import DCGAN_Discriminator, DCGAN_Generator 12 | 13 | 14 | def load_dataset(batch_size=10, download=True): 15 | """ 16 | The output of torchvision datasets are PILImage images of range [0, 1]. 17 | Transform them to Tensors of normalized range [-1, 1] 18 | """ 19 | transform = transforms.Compose([transforms.ToTensor(), 20 | transforms.Normalize((0.5, 0.5, 0.5), 21 | (0.5, 0.5, 0.5))]) 22 | trainset = torchvision.datasets.MNIST(root='../data', train=True, 23 | download=download, 24 | transform=transform) 25 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, 26 | shuffle=True, num_workers=2) 27 | 28 | testset = torchvision.datasets.MNIST(root='../data', train=False, 29 | download=download, 30 | transform=transform) 31 | testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, 32 | shuffle=False, num_workers=2) 33 | 34 | return trainloader, testloader 35 | 36 | 37 | def gen_noise(n_instance, n_dim=2): 38 | """generate n-dim uniform random noise""" 39 | return torch.Tensor(np.random.uniform(low=-1.0, high=1.0, 40 | size=(n_instance, n_dim))) 41 | 42 | 43 | def train_DCGAN(Dis_model, Gen_model, D_criterion, G_criterion, D_optimizer, 44 | G_optimizer, trainloader, n_epoch, batch_size, noise_dim, 45 | n_update_dis=1, n_update_gen=1, use_gpu=False, print_every=10, 46 | update_max=None): 47 | """train DCGAN and print out the losses for D and G""" 48 | for epoch in range(n_epoch): 49 | 50 | D_running_loss = 0.0 51 | G_running_loss = 0.0 52 | 53 | for i, data in enumerate(trainloader, 0): 54 | # get the inputs from true distribution 55 | true_inputs, _ = data 56 | if use_gpu: 57 | true_inputs = true_inputs.cuda() 58 | true_inputs = Variable(true_inputs) 59 | 60 | # get the inputs from the generator 61 | noises = gen_noise(batch_size, n_dim=noise_dim) 62 | if use_gpu: 63 | noises = noises.cuda() 64 | fake_inputs = Gen_model(Variable(noises)) 65 | inputs = torch.cat([true_inputs, fake_inputs]) 66 | 67 | # get the labels 68 | labels = np.zeros(2 * batch_size) 69 | labels[:batch_size] = 1 70 | labels = torch.from_numpy(labels.astype(np.float32)) 71 | if use_gpu: 72 | labels = labels.cuda() 73 | labels = Variable(labels) 74 | 75 | # Discriminator 76 | D_optimizer.zero_grad() 77 | outputs = Dis_model(inputs) 78 | D_loss = D_criterion(outputs[:, 0], labels) 79 | if i % n_update_dis == 0: 80 | D_loss.backward(retain_variables=True) 81 | D_optimizer.step() 82 | 83 | # Generator 84 | if i % n_update_gen == 0: 85 | G_optimizer.zero_grad() 86 | G_loss = G_criterion(outputs[batch_size:, 0], 87 | labels[:batch_size]) 88 | G_loss.backward() 89 | G_optimizer.step() 90 | 91 | # print statistics 92 | D_running_loss += D_loss.data[0] 93 | G_running_loss += G_loss.data[0] 94 | if i % print_every == (print_every - 1): 95 | print('[%d, %5d] D loss: %.3f ; G loss: %.3f' % 96 | (epoch+1, i+1, D_running_loss / print_every, 97 | G_running_loss / print_every)) 98 | D_running_loss = 0.0 99 | G_running_loss = 0.0 100 | 101 | if update_max and i > update_max: 102 | break 103 | 104 | print('Finished Training') 105 | 106 | 107 | def run_DCGAN(n_epoch=2, batch_size=50, use_gpu=False, dis_lr=1e-5, 108 | gen_lr=1e-4, n_update_dis=1, n_update_gen=1, noise_dim=10, 109 | D_featmap_dim=512, G_featmap_dim=1024, n_channel=1, 110 | update_max=None): 111 | # loading data 112 | trainloader, testloader = load_dataset(batch_size=batch_size) 113 | 114 | # initialize models 115 | Dis_model = DCGAN_Discriminator(featmap_dim=D_featmap_dim, 116 | n_channel=n_channel) 117 | Gen_model = DCGAN_Generator(featmap_dim=G_featmap_dim, n_channel=n_channel, 118 | noise_dim=noise_dim) 119 | 120 | if use_gpu: 121 | Dis_model = Dis_model.cuda() 122 | Gen_model = Gen_model.cuda() 123 | 124 | # assign loss function and optimizer (Adam) to D and G 125 | D_criterion = torch.nn.BCELoss() 126 | D_optimizer = optim.Adam(Dis_model.parameters(), lr=dis_lr, 127 | betas=(0.5, 0.999)) 128 | 129 | G_criterion = torch.nn.BCELoss() 130 | G_optimizer = optim.Adam(Gen_model.parameters(), lr=gen_lr, 131 | betas=(0.5, 0.999)) 132 | 133 | train_DCGAN(Dis_model, Gen_model, D_criterion, G_criterion, D_optimizer, 134 | G_optimizer, trainloader, n_epoch, batch_size, noise_dim, 135 | n_update_dis, n_update_gen, update_max=update_max) 136 | 137 | 138 | if __name__ == '__main__': 139 | run_DCGAN(D_featmap_dim=64, G_featmap_dim=128) 140 | -------------------------------------------------------------------------------- /GAN/GAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | class Discriminator(nn.Module): 9 | 10 | def __init__(self): 11 | super(Discriminator, self).__init__() 12 | self.fc1 = nn.Linear(1 * 28 * 28, 256) 13 | self.drop1 = nn.Dropout(p=0.2) 14 | self.fc2 = nn.Linear(256, 128) 15 | self.drop2 = nn.Dropout(p=0.2) 16 | self.fc3 = nn.Linear(128, 1) 17 | 18 | def forward(self, x): 19 | x = self.drop1(F.leaky_relu(self.fc1(x))) 20 | x = self.drop2(F.leaky_relu(self.fc2(x))) 21 | x = F.sigmoid(self.fc3(x)) 22 | return x 23 | 24 | 25 | class Generator(nn.Module): 26 | 27 | def __init__(self): 28 | super(Generator, self).__init__() 29 | self.fc1 = nn.Linear(2, 128) 30 | self.fc2 = nn.Linear(128, 256) 31 | self.fc3 = nn.Linear(256, 512) 32 | self.fc4 = nn.Linear(512, 1 * 28 * 28) 33 | 34 | def forward(self, x): 35 | x = F.elu(self.fc1(x)) 36 | x = F.elu(self.fc2(x)) 37 | x = F.elu(self.fc3(x)) 38 | x = F.tanh(self.fc4(x)) 39 | return x 40 | -------------------------------------------------------------------------------- /GAN/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /GAN/run_GAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch 5 | import torch.optim as optim 6 | import torchvision 7 | import torchvision.transforms as transforms 8 | import numpy as np 9 | 10 | from torch.autograd import Variable 11 | from GAN import Discriminator, Generator 12 | 13 | 14 | def load_dataset(batch_size=10, download=True): 15 | """ 16 | The output of torchvision datasets are PILImage images of range [0, 1]. 17 | Transform them to Tensors of normalized range [-1, 1] 18 | """ 19 | transform = transforms.Compose([transforms.ToTensor(), 20 | transforms.Normalize((0.5, 0.5, 0.5), 21 | (0.5, 0.5, 0.5))]) 22 | trainset = torchvision.datasets.MNIST(root='../data', train=True, 23 | download=download, 24 | transform=transform) 25 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, 26 | shuffle=True, num_workers=2) 27 | 28 | testset = torchvision.datasets.MNIST(root='../data', train=False, 29 | download=download, 30 | transform=transform) 31 | testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, 32 | shuffle=False, num_workers=2) 33 | 34 | return trainloader, testloader 35 | 36 | 37 | def gen_noise(n_instance): 38 | """generate n-dim uniform random noise""" 39 | return torch.Tensor(np.random.uniform(low=-1.0, high=1.0, 40 | size=(n_instance, 2))) 41 | 42 | 43 | def train_GAN(Dis_model, Gen_model, D_criterion, G_criterion, D_optimizer, 44 | G_optimizer, trainloader, n_epoch, batch_size, 45 | n_update_dis=1, n_update_gen=1, use_gpu=False, print_every=10, 46 | update_max=None): 47 | """train GAN and print out the losses for D and G""" 48 | for epoch in range(n_epoch): 49 | 50 | D_running_loss = 0.0 51 | G_running_loss = 0.0 52 | 53 | for i, data in enumerate(trainloader, 0): 54 | # get the inputs from true distribution 55 | true_inputs, _ = data 56 | true_inputs = true_inputs.view(-1, 1 * 28 * 28) 57 | if use_gpu: 58 | true_inputs = true_inputs.cuda() 59 | true_inputs = Variable(true_inputs) 60 | 61 | # get the inputs from the generator 62 | noises = gen_noise(batch_size) 63 | if use_gpu: 64 | noises = noises.cuda() 65 | fake_inputs = Gen_model(Variable(noises)) 66 | inputs = torch.cat([true_inputs, fake_inputs]) 67 | 68 | # get the labels 69 | labels = np.zeros(2 * batch_size) 70 | labels[:batch_size] = 1 71 | labels = torch.from_numpy(labels.astype(np.float32)) 72 | if use_gpu: 73 | labels = labels.cuda() 74 | labels = Variable(labels) 75 | 76 | # Discriminator 77 | D_optimizer.zero_grad() 78 | outputs = Dis_model(inputs) 79 | D_loss = D_criterion(outputs[:, 0], labels) 80 | if i % n_update_dis == 0: 81 | D_loss.backward(retain_variables=True) 82 | D_optimizer.step() 83 | 84 | # Generator 85 | if i % n_update_gen == 0: 86 | G_optimizer.zero_grad() 87 | G_loss = G_criterion(outputs[batch_size:, 0], 88 | labels[:batch_size]) 89 | G_loss.backward() 90 | G_optimizer.step() 91 | 92 | # print statistics 93 | D_running_loss += D_loss.data[0] 94 | G_running_loss += G_loss.data[0] 95 | if i % print_every == (print_every - 1): 96 | print('[%d, %5d] D loss: %.3f ; G loss: %.3f' % 97 | (epoch+1, i+1, D_running_loss / print_every, 98 | G_running_loss / print_every)) 99 | D_running_loss = 0.0 100 | G_running_loss = 0.0 101 | 102 | if update_max and i > update_max: 103 | break 104 | 105 | print('Finished Training') 106 | 107 | 108 | def run_GAN(n_epoch=2, batch_size=50, use_gpu=False, dis_lr=1e-4, gen_lr=1e-3, 109 | n_update_dis=1, n_update_gen=1, update_max=None): 110 | # loading data 111 | trainloader, testloader = load_dataset(batch_size=batch_size) 112 | 113 | # initialize models 114 | Dis_model = Discriminator() 115 | Gen_model = Generator() 116 | 117 | if use_gpu: 118 | Dis_model = Dis_model.cuda() 119 | Gen_model = Gen_model.cuda() 120 | 121 | # assign loss function and optimizer to D and G 122 | D_criterion = torch.nn.BCELoss() 123 | D_optimizer = optim.SGD(Dis_model.parameters(), lr=dis_lr, momentum=0.9) 124 | 125 | G_criterion = torch.nn.BCELoss() 126 | G_optimizer = optim.SGD(Gen_model.parameters(), lr=gen_lr, momentum=0.9) 127 | 128 | train_GAN(Dis_model, Gen_model, D_criterion, G_criterion, D_optimizer, 129 | G_optimizer, trainloader, n_epoch, batch_size, n_update_dis, 130 | n_update_gen, update_max=update_max) 131 | 132 | 133 | if __name__ == '__main__': 134 | run_GAN() 135 | -------------------------------------------------------------------------------- /ImprovedGAN/ImprovedGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | 9 | class ImprovedGAN_Discriminator(nn.Module): 10 | 11 | def __init__(self, featmap_dim=512, n_channel=1, use_gpu=False, 12 | n_B=128, n_C=16): 13 | """ 14 | Minibatch discrimination: learn a tensor to encode side information 15 | from other examples in the same minibatch. 16 | """ 17 | super(ImprovedGAN_Discriminator, self).__init__() 18 | self.use_gpu = use_gpu 19 | self.n_B = n_B 20 | self.n_C = n_C 21 | self.featmap_dim = featmap_dim 22 | 23 | self.conv1 = nn.Conv2d(n_channel, featmap_dim / 4, 5, 24 | stride=2, padding=2) 25 | 26 | self.conv2 = nn.Conv2d(featmap_dim / 4, featmap_dim / 2, 5, 27 | stride=2, padding=2) 28 | self.BN2 = nn.BatchNorm2d(featmap_dim / 2) 29 | 30 | self.conv3 = nn.Conv2d(featmap_dim / 2, featmap_dim, 5, 31 | stride=2, padding=2) 32 | self.BN3 = nn.BatchNorm2d(featmap_dim) 33 | 34 | T_ten_init = torch.randn(featmap_dim * 4 * 4, n_B * n_C) * 0.1 35 | self.T_tensor = nn.Parameter(T_ten_init, requires_grad=True) 36 | self.fc = nn.Linear(featmap_dim * 4 * 4 + n_B, 1) 37 | 38 | def forward(self, x): 39 | """ 40 | Architecture is similar to DCGANs 41 | Add minibatch discrimination => Improved GAN. 42 | """ 43 | x = F.leaky_relu(self.conv1(x), negative_slope=0.2) 44 | x = F.leaky_relu(self.BN2(self.conv2(x)), negative_slope=0.2) 45 | x = F.leaky_relu(self.BN3(self.conv3(x)), negative_slope=0.2) 46 | x = x.view(-1, self.featmap_dim * 4 * 4) 47 | 48 | # #### Minibatch Discrimination ### 49 | T_tensor = self.T_tensor 50 | if self.use_gpu: 51 | T_tensor = T_tensor.cuda() 52 | 53 | Ms = x.mm(T_tensor) 54 | Ms = Ms.view(-1, self.n_B, self.n_C) 55 | 56 | out_tensor = [] 57 | for i in range(Ms.size()[0]): 58 | 59 | out_i = None 60 | for j in range(Ms.size()[0]): 61 | o_i = torch.sum(torch.abs(Ms[i, :, :] - Ms[j, :, :]), 1) 62 | o_i = torch.exp(-o_i) 63 | if out_i is None: 64 | out_i = o_i 65 | else: 66 | out_i = out_i + o_i 67 | 68 | out_tensor.append(out_i) 69 | 70 | out_T = torch.cat(tuple(out_tensor)).view(Ms.size()[0], self.n_B) 71 | x = torch.cat((x, out_T), 1) 72 | # #### Minibatch Discrimination ### 73 | 74 | x = F.sigmoid(self.fc(x)) 75 | 76 | return x 77 | 78 | 79 | class ImprovedGAN_Generator(nn.Module): 80 | 81 | def __init__(self, featmap_dim=1024, n_channel=1, noise_dim=100): 82 | super(ImprovedGAN_Generator, self).__init__() 83 | self.featmap_dim = featmap_dim 84 | self.fc1 = nn.Linear(noise_dim, 4 * 4 * featmap_dim) 85 | self.conv1 = nn.ConvTranspose2d(featmap_dim, (featmap_dim / 2), 5, 86 | stride=2, padding=2) 87 | 88 | self.BN1 = nn.BatchNorm2d(featmap_dim / 2) 89 | self.conv2 = nn.ConvTranspose2d(featmap_dim / 2, featmap_dim / 4, 6, 90 | stride=2, padding=2) 91 | 92 | self.BN2 = nn.BatchNorm2d(featmap_dim / 4) 93 | self.conv3 = nn.ConvTranspose2d(featmap_dim / 4, n_channel, 6, 94 | stride=2, padding=2) 95 | 96 | def forward(self, x): 97 | """ 98 | Project noise to featureMap * width * height, 99 | Batch Normalization after convulation but not at output layer, 100 | ReLU activation function. 101 | """ 102 | x = self.fc1(x) 103 | x = x.view(-1, self.featmap_dim, 4, 4) 104 | x = F.relu(self.BN1(self.conv1(x))) 105 | x = F.relu(self.BN2(self.conv2(x))) 106 | x = F.tanh(self.conv3(x)) 107 | 108 | return x 109 | -------------------------------------------------------------------------------- /ImprovedGAN/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /ImprovedGAN/run_ImprovedGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch 5 | import torch.optim as optim 6 | import torchvision 7 | import torchvision.transforms as transforms 8 | import numpy as np 9 | 10 | from torch.autograd import Variable 11 | from ImprovedGAN import ImprovedGAN_Discriminator, ImprovedGAN_Generator 12 | 13 | 14 | def load_dataset(batch_size=10, download=True): 15 | """ 16 | The output of torchvision datasets are PILImage images of range [0, 1]. 17 | Transform them to Tensors of normalized range [-1, 1] 18 | """ 19 | transform = transforms.Compose([transforms.ToTensor(), 20 | transforms.Normalize((0.5, 0.5, 0.5), 21 | (0.5, 0.5, 0.5))]) 22 | trainset = torchvision.datasets.MNIST(root='../data', train=True, 23 | download=download, 24 | transform=transform) 25 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, 26 | shuffle=True, num_workers=2) 27 | 28 | testset = torchvision.datasets.MNIST(root='../data', train=False, 29 | download=download, 30 | transform=transform) 31 | testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, 32 | shuffle=False, num_workers=2) 33 | 34 | return trainloader, testloader 35 | 36 | 37 | def gen_noise(n_instance, n_dim=2): 38 | """generate n-dim uniform random noise""" 39 | return torch.Tensor(np.random.uniform(low=-1.0, high=1.0, 40 | size=(n_instance, n_dim))) 41 | 42 | 43 | def train_ImprovedGAN(Dis_model, Gen_model, D_criterion, G_criterion, 44 | D_optimizer, G_optimizer, trainloader, n_epoch, 45 | batch_size, noise_dim, n_update_dis=1, n_update_gen=1, 46 | use_gpu=False, print_every=10, update_max=None): 47 | """train ImprovedGAN and print out the losses for D and G""" 48 | for epoch in range(n_epoch): 49 | 50 | D_running_loss = 0.0 51 | G_running_loss = 0.0 52 | 53 | for i, data in enumerate(trainloader, 0): 54 | # get the inputs from true distribution 55 | true_inputs, _ = data 56 | if use_gpu: 57 | true_inputs = true_inputs.cuda() 58 | true_inputs = Variable(true_inputs) 59 | 60 | # get the inputs from the generator 61 | noises = gen_noise(batch_size, n_dim=noise_dim) 62 | if use_gpu: 63 | noises = noises.cuda() 64 | fake_inputs = Gen_model(Variable(noises)) 65 | inputs = torch.cat([true_inputs, fake_inputs]) 66 | 67 | # get the labels 68 | labels = np.zeros(2 * batch_size) 69 | labels[:batch_size] = 1 70 | labels = torch.from_numpy(labels.astype(np.float32)) 71 | if use_gpu: 72 | labels = labels.cuda() 73 | labels = Variable(labels) 74 | 75 | # Discriminator 76 | D_optimizer.zero_grad() 77 | outputs = Dis_model(inputs) 78 | D_loss = D_criterion(outputs[:, 0], labels) 79 | if i % n_update_dis == 0: 80 | D_loss.backward(retain_variables=True) 81 | D_optimizer.step() 82 | 83 | # Generator 84 | if i % n_update_gen == 0: 85 | G_optimizer.zero_grad() 86 | G_loss = G_criterion(outputs[batch_size:, 0], 87 | labels[:batch_size]) 88 | G_loss.backward() 89 | G_optimizer.step() 90 | 91 | # print statistics 92 | D_running_loss += D_loss.data[0] 93 | G_running_loss += G_loss.data[0] 94 | if i % print_every == (print_every - 1): 95 | print('[%d, %5d] D loss: %.3f ; G loss: %.3f' % 96 | (epoch+1, i+1, D_running_loss / print_every, 97 | G_running_loss / print_every)) 98 | D_running_loss = 0.0 99 | G_running_loss = 0.0 100 | 101 | if update_max and i > update_max: 102 | break 103 | 104 | print('Finished Training') 105 | 106 | 107 | def run_ImprovedGAN(n_epoch=2, batch_size=50, use_gpu=False, dis_lr=1e-5, 108 | gen_lr=1e-4, n_update_dis=1, n_update_gen=1, noise_dim=10, 109 | D_featmap_dim=512, G_featmap_dim=1024, n_channel=1, 110 | update_max=None): 111 | # loading data 112 | trainloader, testloader = load_dataset(batch_size=batch_size) 113 | 114 | # initialize models 115 | Dis_model = ImprovedGAN_Discriminator(featmap_dim=D_featmap_dim, 116 | n_channel=n_channel) 117 | Gen_model = ImprovedGAN_Generator(featmap_dim=G_featmap_dim, 118 | n_channel=n_channel, 119 | noise_dim=noise_dim) 120 | 121 | if use_gpu: 122 | Dis_model = Dis_model.cuda() 123 | Gen_model = Gen_model.cuda() 124 | 125 | # assign loss function and optimizer (Adam) to D and G 126 | D_criterion = torch.nn.BCELoss() 127 | D_optimizer = optim.Adam(Dis_model.parameters(), lr=dis_lr, 128 | betas=(0.5, 0.999)) 129 | 130 | G_criterion = torch.nn.BCELoss() 131 | G_optimizer = optim.Adam(Gen_model.parameters(), lr=gen_lr, 132 | betas=(0.5, 0.999)) 133 | 134 | train_ImprovedGAN(Dis_model, Gen_model, D_criterion, G_criterion, 135 | D_optimizer, G_optimizer, trainloader, n_epoch, 136 | batch_size, noise_dim, n_update_dis, n_update_gen, 137 | update_max=update_max) 138 | 139 | 140 | if __name__ == '__main__': 141 | run_ImprovedGAN(D_featmap_dim=64, G_featmap_dim=128, batch_size=10) 142 | -------------------------------------------------------------------------------- /InfoGAN/InfoGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | 7 | 8 | class InfoGAN_Discriminator(nn.Module): 9 | 10 | def __init__(self, n_layer=3, n_conti=2, n_discrete=1, 11 | num_category=10, use_gpu=False, featmap_dim=256, 12 | n_channel=1): 13 | """ 14 | InfoGAN Discriminator, have additional outputs for latent codes. 15 | Architecture brought from DCGAN. 16 | """ 17 | super(InfoGAN_Discriminator, self).__init__() 18 | self.n_layer = n_layer 19 | self.n_conti = n_conti 20 | self.n_discrete = n_discrete 21 | self.num_category = num_category 22 | 23 | # Discriminator 24 | self.featmap_dim = featmap_dim 25 | 26 | convs = [] 27 | BNs = [] 28 | for layer in range(self.n_layer): 29 | if layer == (self.n_layer - 1): 30 | n_conv_in = n_channel 31 | else: 32 | n_conv_in = int(featmap_dim / (2**(layer + 1))) 33 | n_conv_out = int(featmap_dim / (2**layer)) 34 | 35 | _conv = nn.Conv2d(n_conv_in, n_conv_out, kernel_size=5, 36 | stride=2, padding=2) 37 | if use_gpu: 38 | _conv = _conv.cuda() 39 | convs.append(_conv) 40 | 41 | if layer != (self.n_layer - 1): 42 | _BN = nn.BatchNorm2d(n_conv_out) 43 | if use_gpu: 44 | _BN = _BN.cuda() 45 | BNs.append(_BN) 46 | 47 | # output layer - prob(real) and auxiliary distributions Q(c_j|x) 48 | n_hidden = featmap_dim * 4 * 4 49 | n_output = 1 + n_conti + n_discrete * num_category 50 | self.fc = nn.Linear(n_hidden, n_output) 51 | 52 | # register all nn modules 53 | self.convs = nn.ModuleList(convs) 54 | self.BNs = nn.ModuleList(BNs) 55 | 56 | def forward(self, x): 57 | """ 58 | Output the probability of being in real dataset 59 | plus the conditional distributions of latent codes. 60 | """ 61 | for layer in range(self.n_layer): 62 | conv_layer = self.convs[self.n_layer - layer - 1] 63 | 64 | if layer == 0: 65 | x = F.leaky_relu(conv_layer(x), negative_slope=0.2) 66 | else: 67 | BN_layer = self.BNs[self.n_layer - layer - 1] 68 | x = F.leaky_relu(BN_layer(conv_layer(x)), negative_slope=0.2) 69 | 70 | x = x.view(-1, self.featmap_dim * 4 * 4) 71 | 72 | # output layer 73 | x = self.fc(x) 74 | x[:, 0] = F.sigmoid(x[:, 0].clone()) 75 | for j in range(self.n_discrete): 76 | start = 1 + self.n_conti + j * self.num_category 77 | end = start + self.num_category 78 | x[:, start:end] = F.softmax(x[:, start:end].clone()) 79 | 80 | return x 81 | 82 | 83 | class InfoGAN_Generator(nn.Module): 84 | 85 | def __init__(self, noise_dim=10, n_layer=3, n_conti=2, n_discrete=1, 86 | num_category=10, use_gpu=False, featmap_dim=256, n_channel=1): 87 | """ 88 | InfoGAN Generator, have an additional input branch for latent codes. 89 | Architecture brought from DCGAN. 90 | """ 91 | super(InfoGAN_Generator, self).__init__() 92 | self.n_layer = n_layer 93 | self.n_conti = n_conti 94 | self.n_discrete = n_discrete 95 | self.num_category = num_category 96 | 97 | # calculate input dimension 98 | n_input = noise_dim + n_conti + n_discrete * num_category 99 | 100 | # Generator 101 | self.featmap_dim = featmap_dim 102 | self.fc_in = nn.Linear(n_input, featmap_dim * 4 * 4) 103 | 104 | convs = [] 105 | BNs = [] 106 | for layer in range(self.n_layer): 107 | if layer == 0: 108 | n_conv_out = n_channel 109 | else: 110 | n_conv_out = featmap_dim / (2 ** (self.n_layer - layer)) 111 | 112 | n_conv_in = featmap_dim / (2 ** (self.n_layer - layer - 1)) 113 | n_width = 5 if layer == (self.n_layer - 1) else 6 114 | 115 | _conv = nn.ConvTranspose2d(n_conv_in, n_conv_out, n_width, 116 | stride=2, padding=2) 117 | 118 | if use_gpu: 119 | _conv = _conv.cuda() 120 | convs.append(_conv) 121 | 122 | if layer != 0: 123 | _BN = nn.BatchNorm2d(n_conv_out) 124 | if use_gpu: 125 | _BN = _BN.cuda() 126 | BNs.append(_BN) 127 | 128 | # register all nn modules 129 | self.convs = nn.ModuleList(convs) 130 | self.BNs = nn.ModuleList(BNs) 131 | 132 | def forward(self, x): 133 | """ 134 | Input the random noise plus latent codes to generate fake images. 135 | """ 136 | x = self.fc_in(x) 137 | x = x.view(-1, self.featmap_dim, 4, 4) 138 | 139 | for layer in range(self.n_layer): 140 | conv_layer = self.convs[self.n_layer - layer - 1] 141 | if layer == (self.n_layer - 1): 142 | x = F.tanh(conv_layer(x)) 143 | else: 144 | BN_layer = self.BNs[self.n_layer - layer - 2] 145 | x = F.relu(BN_layer(conv_layer(x))) 146 | 147 | return x 148 | -------------------------------------------------------------------------------- /InfoGAN/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /InfoGAN/run_InfoGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch 5 | import torch.optim as optim 6 | import torchvision 7 | import torchvision.transforms as transforms 8 | import numpy as np 9 | 10 | from torch.autograd import Variable 11 | from InfoGAN import InfoGAN_Discriminator, InfoGAN_Generator 12 | 13 | 14 | def load_dataset(batch_size=10, download=True): 15 | """ 16 | The output of torchvision datasets are PILImage images of range [0, 1]. 17 | Transform them to Tensors of normalized range [-1, 1] 18 | """ 19 | transform = transforms.Compose([transforms.ToTensor(), 20 | transforms.Normalize((0.5, 0.5, 0.5), 21 | (0.5, 0.5, 0.5))]) 22 | trainset = torchvision.datasets.MNIST(root='../data', train=True, 23 | download=download, 24 | transform=transform) 25 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, 26 | shuffle=True, num_workers=2) 27 | 28 | testset = torchvision.datasets.MNIST(root='../data', train=False, 29 | download=download, 30 | transform=transform) 31 | testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, 32 | shuffle=False, num_workers=2) 33 | 34 | return trainloader, testloader 35 | 36 | 37 | def gen_noise(n_instance, n_dim=2): 38 | """generate n-dim uniform random noise""" 39 | return torch.Tensor(np.random.uniform(low=-1.0, high=1.0, 40 | size=(n_instance, n_dim))) 41 | 42 | 43 | def gen_conti_codes(n_instance, n_conti, mean=0, std=1): 44 | """generate gaussian continuous codes with specified mean and std""" 45 | codes = np.random.randn(n_instance, n_conti) * std + mean 46 | return torch.Tensor(codes) 47 | 48 | 49 | def gen_discrete_code(n_instance, n_discrete, num_category=10): 50 | """generate discrete codes with n categories""" 51 | codes = [] 52 | for i in range(n_discrete): 53 | code = np.zeros((n_instance, num_category)) 54 | random_cate = np.random.randint(0, num_category, n_instance) 55 | code[range(n_instance), random_cate] = 1 56 | codes.append(code) 57 | 58 | codes = np.concatenate(codes, 1) 59 | return torch.Tensor(codes) 60 | 61 | 62 | def train_InfoGAN(InfoGAN_Dis, InfoGAN_Gen, D_criterion, G_criterion, 63 | D_optimizer, G_optimizer, info_reg_discrete, info_reg_conti, 64 | n_conti, n_discrete, mean, std, num_category, trainloader, 65 | n_epoch, batch_size, noise_dim, 66 | n_update_dis=1, n_update_gen=1, use_gpu=False, 67 | print_every=50, update_max=None): 68 | """train InfoGAN and print out the losses for D and G""" 69 | 70 | for epoch in range(n_epoch): 71 | 72 | D_running_loss = 0.0 73 | G_running_loss = 0.0 74 | 75 | for i, data in enumerate(trainloader, 0): 76 | # get the inputs from true distribution 77 | true_inputs, lab = data 78 | true_inputs = Variable(true_inputs) 79 | if use_gpu: 80 | true_inputs = true_inputs.cuda() 81 | 82 | # get inputs (noises and codes) for Generator 83 | noises = Variable(gen_noise(batch_size, n_dim=noise_dim)) 84 | conti_codes = Variable(gen_conti_codes(batch_size, n_conti, 85 | mean, std)) 86 | discr_codes = Variable(gen_discrete_code(batch_size, n_discrete, 87 | num_category)) 88 | if use_gpu: 89 | noises = noises.cuda() 90 | conti_codes = conti_codes.cuda() 91 | discr_codes = discr_codes.cuda() 92 | 93 | # generate fake images 94 | gen_inputs = torch.cat((noises, conti_codes, discr_codes), 1) 95 | fake_inputs = InfoGAN_Gen(gen_inputs) 96 | inputs = torch.cat([true_inputs, fake_inputs]) 97 | 98 | # make a minibatch of labels 99 | labels = np.zeros(2 * batch_size) 100 | labels[:batch_size] = 1 101 | labels = torch.from_numpy(labels.astype(np.float32)) 102 | if use_gpu: 103 | labels = labels.cuda() 104 | labels = Variable(labels) 105 | 106 | # Discriminator 107 | D_optimizer.zero_grad() 108 | outputs = InfoGAN_Dis(inputs) 109 | 110 | # calculate mutual information lower bound L(G, Q) 111 | for j in range(n_discrete): 112 | shift = (j * num_category) 113 | start = 1 + n_conti + shift 114 | end = start + num_category 115 | Q_cx_discr = outputs[batch_size:, start:end] 116 | codes = discr_codes[:, shift:(shift+num_category)] 117 | condi_entro = -torch.mean(torch.sum(Q_cx_discr * codes, 1)) 118 | 119 | if j == 0: 120 | L_discrete = -condi_entro 121 | else: 122 | L_discrete -= condi_entro 123 | L_discrete /= n_discrete 124 | 125 | Q_cx_conti = outputs[batch_size:, 1:(1 + n_conti)] 126 | L_conti = torch.mean(-(((Q_cx_conti - mean) / std) ** 2)) 127 | 128 | # Update Discriminator 129 | D_loss = D_criterion(outputs[:, 0], labels) 130 | if n_discrete > 0: 131 | D_loss = D_loss - info_reg_discrete * L_discrete 132 | 133 | if n_conti > 0: 134 | D_loss = D_loss - info_reg_conti * L_conti 135 | 136 | if i % n_update_dis == 0: 137 | D_loss.backward(retain_variables=True) 138 | D_optimizer.step() 139 | 140 | # Update Generator 141 | if i % n_update_gen == 0: 142 | G_optimizer.zero_grad() 143 | G_loss = G_criterion(outputs[batch_size:, 0], 144 | labels[:batch_size]) 145 | 146 | if n_discrete > 0: 147 | G_loss = G_loss - info_reg_discrete * L_discrete 148 | 149 | if n_conti > 0: 150 | G_loss = G_loss - info_reg_conti * L_conti 151 | 152 | G_loss.backward() 153 | G_optimizer.step() 154 | 155 | # print statistics 156 | D_running_loss += D_loss.data[0] 157 | G_running_loss += G_loss.data[0] 158 | if i % print_every == (print_every - 1): 159 | print('[%d, %5d] D loss: %.3f ; G loss: %.3f' % 160 | (epoch+1, i+1, D_running_loss / print_every, 161 | G_running_loss / print_every)) 162 | D_running_loss = 0.0 163 | G_running_loss = 0.0 164 | 165 | if update_max and i > update_max: 166 | break 167 | 168 | print('Finished Training') 169 | 170 | 171 | def run_InfoGAN(info_reg_discrete=1.0, info_reg_conti=0.5, noise_dim=10, 172 | n_conti=2, n_discrete=1, mean=0.0, std=0.5, num_category=10, 173 | n_layer=3, n_channel=1, D_featmap_dim=256, G_featmap_dim=1024, 174 | n_epoch=2, batch_size=50, use_gpu=False, dis_lr=1e-4, 175 | gen_lr=1e-3, n_update_dis=1, n_update_gen=1, update_max=None): 176 | # loading data 177 | trainloader, testloader = load_dataset(batch_size=batch_size) 178 | 179 | # initialize models 180 | InfoGAN_Dis = InfoGAN_Discriminator(n_layer, n_conti, n_discrete, 181 | num_category, use_gpu, D_featmap_dim, 182 | n_channel) 183 | 184 | InfoGAN_Gen = InfoGAN_Generator(noise_dim, n_layer, n_conti, n_discrete, 185 | num_category, use_gpu, G_featmap_dim, 186 | n_channel) 187 | 188 | if use_gpu: 189 | InfoGAN_Dis = InfoGAN_Dis.cuda() 190 | InfoGAN_Gen = InfoGAN_Gen.cuda() 191 | 192 | # assign loss function and optimizer (Adam) to D and G 193 | D_criterion = torch.nn.BCELoss() 194 | D_optimizer = optim.Adam(InfoGAN_Dis.parameters(), lr=dis_lr, 195 | betas=(0.5, 0.999)) 196 | 197 | G_criterion = torch.nn.BCELoss() 198 | G_optimizer = optim.Adam(InfoGAN_Gen.parameters(), lr=gen_lr, 199 | betas=(0.5, 0.999)) 200 | 201 | train_InfoGAN(InfoGAN_Dis, InfoGAN_Gen, D_criterion, G_criterion, 202 | D_optimizer, G_optimizer, info_reg_discrete, info_reg_conti, 203 | n_conti, n_discrete, mean, std, num_category, trainloader, 204 | n_epoch, batch_size, noise_dim, 205 | n_update_dis, n_update_gen, use_gpu, update_max=update_max) 206 | 207 | 208 | if __name__ == '__main__': 209 | run_InfoGAN(n_conti=2, n_discrete=1, D_featmap_dim=64, G_featmap_dim=128, 210 | n_epoch=1, batch_size=10, update_max=200) 211 | -------------------------------------------------------------------------------- /LAPGAN/LAPGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | import numpy as np 8 | import cv2 9 | 10 | from torch.autograd import Variable 11 | 12 | 13 | class CondiGAN_Discriminator(nn.Module): 14 | 15 | def __init__(self, n_layer=3, condition=True, n_condition=100, 16 | use_gpu=False, featmap_dim=256, n_channel=1, 17 | condi_featmap_dim=256): 18 | """ 19 | Conditional Discriminator. 20 | Architecture brought from DCGAN. 21 | """ 22 | super(CondiGAN_Discriminator, self).__init__() 23 | self.n_layer = n_layer 24 | self.condition = condition 25 | 26 | # original Discriminator 27 | self.featmap_dim = featmap_dim 28 | 29 | convs = [] 30 | BNs = [] 31 | for layer in range(self.n_layer): 32 | if layer == (self.n_layer - 1): 33 | n_conv_in = n_channel 34 | else: 35 | n_conv_in = int(featmap_dim / (2**(layer + 1))) 36 | n_conv_out = int(featmap_dim / (2**layer)) 37 | 38 | _conv = nn.Conv2d(n_conv_in, n_conv_out, kernel_size=5, 39 | stride=2, padding=2) 40 | if use_gpu: 41 | _conv = _conv.cuda() 42 | convs.append(_conv) 43 | 44 | if layer != (self.n_layer - 1): 45 | _BN = nn.BatchNorm2d(n_conv_out) 46 | if use_gpu: 47 | _BN = _BN.cuda() 48 | BNs.append(_BN) 49 | 50 | # extra image information to be conditioned on 51 | if self.condition: 52 | self.condi_featmap_dim = condi_featmap_dim 53 | convs_condi = [] 54 | BNs_condi = [] 55 | 56 | for layer in range(self.n_layer): 57 | if layer == (self.n_layer - 1): 58 | n_conv_in = n_channel 59 | else: 60 | n_conv_in = int(condi_featmap_dim / (2**(layer + 1))) 61 | n_conv_out = int(condi_featmap_dim / (2**layer)) 62 | 63 | _conv = nn.Conv2d(n_conv_in, n_conv_out, kernel_size=5, 64 | stride=2, padding=2) 65 | if use_gpu: 66 | _conv = _conv.cuda() 67 | convs_condi.append(_conv) 68 | 69 | if layer != (self.n_layer - 1): 70 | _BN = nn.BatchNorm2d(n_conv_out) 71 | if use_gpu: 72 | _BN = _BN.cuda() 73 | BNs_condi.append(_BN) 74 | 75 | self.fc_c = nn.Linear(condi_featmap_dim * 4 * 4, n_condition) 76 | 77 | # register layer modules 78 | self.convs = nn.ModuleList(convs) 79 | self.BNs = nn.ModuleList(BNs) 80 | if self.condition: 81 | self.convs_condi = nn.ModuleList(convs_condi) 82 | self.BNs_condi = nn.ModuleList(BNs_condi) 83 | 84 | # output layer 85 | n_hidden = featmap_dim * 4 * 4 86 | if self.condition: 87 | n_hidden += n_condition 88 | self.fc = nn.Linear(n_hidden, 1) 89 | 90 | def forward(self, x, condi_x=None): 91 | """ 92 | Concatenate CNN-processed extra information vector at the last layer 93 | """ 94 | for layer in range(self.n_layer): 95 | conv_layer = self.convs[self.n_layer - layer - 1] 96 | if layer == 0: 97 | x = F.leaky_relu(conv_layer(x), negative_slope=0.2) 98 | else: 99 | BN_layer = self.BNs[self.n_layer - layer - 1] 100 | x = F.leaky_relu(BN_layer(conv_layer(x)), negative_slope=0.2) 101 | x = x.view(-1, self.featmap_dim * 4 * 4) 102 | 103 | # calculate and concatenate extra information 104 | if self.condition: 105 | for layer in range(self.n_layer): 106 | _conv = self.convs_condi[self.n_layer - layer - 1] 107 | if layer == 0: 108 | condi_x = F.leaky_relu(_conv(condi_x), negative_slope=0.2) 109 | else: 110 | BN_layer = self.BNs_condi[self.n_layer - layer - 1] 111 | condi_x = F.leaky_relu(BN_layer(_conv(condi_x)), 112 | negative_slope=0.2) 113 | 114 | condi_x = condi_x.view(-1, self.condi_featmap_dim * 4 * 4) 115 | condi_x = self.fc_c(condi_x) 116 | x = torch.cat((x, condi_x), 1) 117 | 118 | # output layer 119 | x = F.sigmoid(self.fc(x)) 120 | 121 | return x 122 | 123 | 124 | class CondiGAN_Generator(nn.Module): 125 | 126 | def __init__(self, noise_dim=10, n_layer=3, condition=True, 127 | n_condition=100, use_gpu=False, featmap_dim=256, n_channel=1, 128 | condi_featmap_dim=256): 129 | """ 130 | Conditional Generator. 131 | Architecture brought from DCGAN. 132 | """ 133 | super(CondiGAN_Generator, self).__init__() 134 | self.n_layer = n_layer 135 | self.condition = condition 136 | 137 | # extra image information to be conditioned on 138 | if self.condition: 139 | self.condi_featmap_dim = condi_featmap_dim 140 | 141 | convs_condi = [] 142 | BNs_condi = [] 143 | for layer in range(self.n_layer): 144 | if layer == (self.n_layer - 1): 145 | n_conv_in = n_channel 146 | else: 147 | n_conv_in = int(condi_featmap_dim / (2**(layer + 1))) 148 | n_conv_out = int(condi_featmap_dim / (2**layer)) 149 | 150 | _conv = nn.Conv2d(n_conv_in, n_conv_out, kernel_size=5, 151 | stride=2, padding=2) 152 | if use_gpu: 153 | _conv = _conv.cuda() 154 | convs_condi.append(_conv) 155 | 156 | if layer != (self.n_layer - 1): 157 | _BN = nn.BatchNorm2d(n_conv_out) 158 | if use_gpu: 159 | _BN = _BN.cuda() 160 | BNs_condi.append(_BN) 161 | 162 | self.fc_c = nn.Linear(condi_featmap_dim * 4 * 4, n_condition) 163 | 164 | # calculate input dimension 165 | n_input = noise_dim 166 | if self.condition: 167 | n_input += n_condition 168 | 169 | # Generator 170 | self.featmap_dim = featmap_dim 171 | self.fc1 = nn.Linear(n_input, int(featmap_dim * 4 * 4)) 172 | 173 | convs = [] 174 | BNs = [] 175 | for layer in range(self.n_layer): 176 | if layer == 0: 177 | n_conv_out = n_channel 178 | else: 179 | n_conv_out = featmap_dim / (2 ** (self.n_layer - layer)) 180 | n_conv_in = featmap_dim / (2 ** (self.n_layer - layer - 1)) 181 | 182 | n_width = 5 if layer == (self.n_layer - 1) else 6 183 | _conv = nn.ConvTranspose2d(n_conv_in, n_conv_out, n_width, 184 | stride=2, padding=2) 185 | if use_gpu: 186 | _conv = _conv.cuda() 187 | convs.append(_conv) 188 | 189 | if layer != 0: 190 | _BN = nn.BatchNorm2d(n_conv_out) 191 | if use_gpu: 192 | _BN = _BN.cuda() 193 | BNs.append(_BN) 194 | 195 | # register layer modules 196 | self.convs = nn.ModuleList(convs) 197 | self.BNs = nn.ModuleList(BNs) 198 | if self.condition: 199 | self.convs_condi = nn.ModuleList(convs_condi) 200 | self.BNs_condi = nn.ModuleList(BNs_condi) 201 | 202 | def forward(self, x, condi_x=None): 203 | """ 204 | Concatenate CNN-processed extra information vector at the first layer 205 | """ 206 | # calculate and concatenate extra information 207 | if self.condition: 208 | for layer in range(self.n_layer): 209 | _conv = self.convs_condi[self.n_layer - layer - 1] 210 | if layer == 0: 211 | condi_x = F.leaky_relu(_conv(condi_x), negative_slope=0.2) 212 | else: 213 | BN_layer = self.BNs_condi[self.n_layer - layer - 1] 214 | condi_x = F.leaky_relu(BN_layer(_conv(condi_x)), 215 | negative_slope=0.2) 216 | 217 | condi_x = condi_x.view(-1, self.condi_featmap_dim * 4 * 4) 218 | condi_x = self.fc_c(condi_x) 219 | x = torch.cat((x, condi_x), 1) 220 | 221 | x = self.fc1(x) 222 | x = x.view(-1, self.featmap_dim, 4, 4) 223 | 224 | for layer in range(self.n_layer): 225 | conv_layer = self.convs[self.n_layer - layer - 1] 226 | if layer == (self.n_layer - 1): 227 | x = F.tanh(conv_layer(x)) 228 | else: 229 | BN_layer = self.BNs[self.n_layer - layer - 2] 230 | x = F.relu(BN_layer(conv_layer(x))) 231 | 232 | return x 233 | 234 | 235 | class LAPGAN(object): 236 | 237 | def __init__(self, n_level, noise_dim=10, n_condition=100, 238 | D_featmap_dim=64, condi_D_featmap_dim=64, 239 | G_featmap_dim=256, condi_G_featmap_dim=64, 240 | use_gpu=False, n_channel=1): 241 | """ 242 | Initialize a group of discriminators and generators for LAPGAN 243 | n_level: number of levels in the Laplacian Pyramid 244 | noise_dim: dimension of random noise to feed into the last generator 245 | D_featmap_dim: discriminator, (#feature maps) in the last layer of CNN 246 | condi_D_featmap_dim: (#feature maps) of extra info CNN's last layer 247 | G_featmap_dim: generator, (#feature maps) of deconvNN's first layer 248 | condi_G_featmap_dim: (#feature maps) of extra info CNN's last layer 249 | use_gpu: to use GPU computation or not 250 | n_channel: number of channel for input images 251 | """ 252 | self.n_level = n_level 253 | self.n_channel = n_channel 254 | self.use_gpu = use_gpu 255 | self.noise_dim = noise_dim 256 | self.Dis_models = [] 257 | self.Gen_models = [] 258 | 259 | for level in range(n_level): 260 | n_layer = n_level - level 261 | if level == (n_level - 1): 262 | condition = False 263 | else: 264 | condition = True 265 | 266 | Dis_model = CondiGAN_Discriminator(n_layer, condition, n_condition, 267 | use_gpu, D_featmap_dim, 268 | n_channel, condi_D_featmap_dim) 269 | Gen_model = CondiGAN_Generator(noise_dim, n_layer, condition, 270 | n_condition, use_gpu, G_featmap_dim, 271 | n_channel, condi_G_featmap_dim) 272 | 273 | if use_gpu: 274 | Dis_model = Dis_model.cuda() 275 | Gen_model = Gen_model.cuda() 276 | 277 | self.Dis_models.append(Dis_model) 278 | self.Gen_models.append(Gen_model) 279 | 280 | def generate(self, batchsize, get_level=None, generator=False): 281 | """Generate images from LAPGAN generators""" 282 | self.outputs = [] 283 | self.generator_outputs = [] 284 | for level in range(self.n_level): 285 | Gen_model = self.Gen_models[self.n_level - level - 1] 286 | 287 | # generate noise 288 | noise = Variable(gen_noise(batchsize, self.noise_dim)) 289 | if self.use_gpu: 290 | noise = noise.cuda() 291 | 292 | if level == 0: 293 | # directly generate images 294 | output_imgs = Gen_model.forward(noise) 295 | if self.use_gpu: 296 | output_imgs = output_imgs.cpu() 297 | output_imgs = output_imgs.data.numpy() 298 | self.generator_outputs.append(output_imgs) 299 | else: 300 | # upsize 301 | input_imgs = np.array([[cv2.pyrUp(output_imgs[i, j, :]) 302 | for j in range(self.n_channel)] 303 | for i in range(batchsize)]) 304 | condi_imgs = Variable(torch.Tensor(input_imgs)) 305 | if self.use_gpu: 306 | condi_imgs = condi_imgs.cuda() 307 | 308 | # generate images with extra information 309 | residual_imgs = Gen_model.forward(noise, condi_imgs) 310 | if self.use_gpu: 311 | residual_imgs = residual_imgs.cpu() 312 | output_imgs = residual_imgs.data.numpy() + input_imgs 313 | self.generator_outputs.append(residual_imgs.data.numpy()) 314 | 315 | self.outputs.append(output_imgs) 316 | 317 | if get_level is None: 318 | get_level = -1 319 | 320 | if generator: 321 | result_imgs = self.generator_outputs[get_level] 322 | else: 323 | result_imgs = self.outputs[get_level] 324 | 325 | return result_imgs 326 | 327 | 328 | def gen_noise(n_instance, n_dim=2): 329 | """generate n-dim uniform random noise""" 330 | return torch.Tensor(np.random.uniform(low=-1.0, high=1.0, 331 | size=(n_instance, n_dim))) 332 | -------------------------------------------------------------------------------- /LAPGAN/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /LAPGAN/run_LAPGAN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # @Author: aaronlai 3 | 4 | import torch 5 | import torch.nn as nn 6 | import torch.optim as optim 7 | import torchvision 8 | import torchvision.transforms as transforms 9 | import numpy as np 10 | import cv2 11 | 12 | from torch.autograd import Variable 13 | from LAPGAN import LAPGAN, gen_noise 14 | 15 | 16 | def load_dataset(batch_size=10, download=True): 17 | """ 18 | The output of torchvision datasets are PILImage images of range [0, 1]. 19 | Transform them to Tensors of normalized range [-1, 1] 20 | """ 21 | transform = transforms.Compose([transforms.ToTensor(), 22 | transforms.Normalize((0.5, 0.5, 0.5), 23 | (0.5, 0.5, 0.5))]) 24 | trainset = torchvision.datasets.MNIST(root='../data', train=True, 25 | download=download, 26 | transform=transform) 27 | trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, 28 | shuffle=True, num_workers=2) 29 | 30 | testset = torchvision.datasets.MNIST(root='../data', train=False, 31 | download=download, 32 | transform=transform) 33 | testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, 34 | shuffle=False, num_workers=2) 35 | 36 | return trainloader, testloader 37 | 38 | 39 | def train_LAPGAN(LapGan_model, n_level, D_criterions, G_criterions, 40 | D_optimizers, G_optimizers, trainloader, n_epoch, 41 | batch_size, noise_dim, n_update_dis=1, n_update_gen=1, 42 | use_gpu=False, print_every=10, update_max=None): 43 | """train LAPGAN and print out the losses for Ds and Gs""" 44 | for epoch in range(n_epoch): 45 | 46 | D_running_losses = [0.0 for i in range(n_level)] 47 | G_running_losses = [0.0 for i in range(n_level)] 48 | 49 | for ind, data in enumerate(trainloader, 0): 50 | # get the inputs from true distribution 51 | true_inputs, lab = data 52 | down_imgs = true_inputs.numpy() 53 | n_minibatch, n_channel, _, _ = down_imgs.shape 54 | 55 | for l in range(n_level): 56 | # calculate input images for models at the particular level 57 | if l == (n_level - 1): 58 | condi_inputs = None 59 | true_inputs = Variable(torch.Tensor(down_imgs)) 60 | if use_gpu: 61 | true_inputs = true_inputs.cuda() 62 | else: 63 | new_down_imgs = [] 64 | up_imgs = [] 65 | residual_imgs = [] 66 | 67 | # compute a Laplacian Pyramid 68 | for i in range(n_minibatch): 69 | down_img = [] 70 | up_img = [] 71 | residual_img = [] 72 | 73 | for j in range(n_channel): 74 | previous = down_imgs[i, j, :] 75 | down_img.append(cv2.pyrDown(previous)) 76 | up_img.append(cv2.pyrUp(down_img[-1])) 77 | residual_img.append(previous - up_img[-1]) 78 | 79 | new_down_imgs.append(down_img) 80 | up_imgs.append(up_img) 81 | residual_imgs.append(residual_img) 82 | 83 | down_imgs = np.array(new_down_imgs) 84 | up_imgs = np.array(up_imgs) 85 | residual_imgs = np.array(residual_imgs) 86 | 87 | condi_inputs = Variable(torch.Tensor(up_imgs)) 88 | true_inputs = Variable(torch.Tensor(residual_imgs)) 89 | if use_gpu: 90 | condi_inputs = condi_inputs.cuda() 91 | true_inputs = true_inputs.cuda() 92 | 93 | # get inputs for discriminators from generators and real data 94 | noise = Variable(gen_noise(batch_size, noise_dim)) 95 | if use_gpu: 96 | noise = noise.cuda() 97 | fake_inputs = LapGan_model.Gen_models[l](noise, condi_inputs) 98 | inputs = torch.cat([true_inputs, fake_inputs]) 99 | labels = np.zeros(2 * batch_size) 100 | labels[:batch_size] = 1 101 | labels = Variable(torch.from_numpy(labels.astype(np.float32))) 102 | if use_gpu: 103 | labels = labels.cuda() 104 | 105 | # Discriminator 106 | D_optimizers[l].zero_grad() 107 | if condi_inputs: 108 | condi_inputs = torch.cat((condi_inputs, condi_inputs)) 109 | outputs = LapGan_model.Dis_models[l](inputs, condi_inputs) 110 | D_loss = D_criterions[l](outputs[:, 0], labels) 111 | 112 | if ind % n_update_dis == 0: 113 | D_loss.backward(retain_variables=True) 114 | D_optimizers[l].step() 115 | 116 | # Generator 117 | if ind % n_update_gen == 0: 118 | G_optimizers[l].zero_grad() 119 | G_loss = G_criterions[l](outputs[batch_size:, 0], 120 | labels[:batch_size]) 121 | G_loss.backward() 122 | G_optimizers[l].step() 123 | 124 | # print statistics 125 | D_running_losses[l] += D_loss.data[0] 126 | G_running_losses[l] += G_loss.data[0] 127 | if ind % print_every == (print_every - 1): 128 | print('[%d, %5d, %d] D loss: %.3f ; G loss: %.3f' % 129 | (epoch+1, ind+1, l+1, 130 | D_running_losses[l] / print_every, 131 | G_running_losses[l] / print_every)) 132 | D_running_losses[l] = 0.0 133 | G_running_losses[l] = 0.0 134 | 135 | if update_max and ind > update_max: 136 | break 137 | 138 | print('Finished Training') 139 | 140 | 141 | def run_LAPGAN(n_level=3, n_epoch=2, batch_size=20, use_gpu=False, 142 | dis_lrs=None, gen_lrs=None, n_update_dis=1, n_update_gen=1, 143 | noise_dim=10, n_condition=100, D_featmap_dim=128, 144 | condi_D_featmap_dim=128, G_featmap_dim=256, 145 | condi_G_featmap_dim=128, n_channel=1, n_sample=25, 146 | update_max=None): 147 | # loading data 148 | trainloader, testloader = load_dataset(batch_size=batch_size) 149 | 150 | # initialize models 151 | LapGan_model = LAPGAN(n_level, noise_dim, n_condition, D_featmap_dim, 152 | condi_D_featmap_dim, G_featmap_dim, 153 | condi_G_featmap_dim, use_gpu, n_channel) 154 | 155 | # assign loss function and optimizer (Adam) to D and G 156 | D_criterions = [] 157 | G_criterions = [] 158 | 159 | D_optimizers = [] 160 | G_optimizers = [] 161 | 162 | if not dis_lrs: 163 | dis_lrs = [0.0002, 0.0003, 0.001] 164 | 165 | if not gen_lrs: 166 | gen_lrs = [0.001, 0.005, 0.01] 167 | 168 | for l in range(n_level): 169 | D_criterions.append(nn.BCELoss()) 170 | D_optim = optim.Adam(LapGan_model.Dis_models[l].parameters(), 171 | lr=dis_lrs[l], betas=(0.5, 0.999)) 172 | D_optimizers.append(D_optim) 173 | 174 | G_criterions.append(nn.BCELoss()) 175 | G_optim = optim.Adam(LapGan_model.Gen_models[l].parameters(), 176 | lr=gen_lrs[l], betas=(0.5, 0.999)) 177 | G_optimizers.append(G_optim) 178 | 179 | train_LAPGAN(LapGan_model, n_level, D_criterions, G_criterions, 180 | D_optimizers, G_optimizers, trainloader, n_epoch, 181 | batch_size, noise_dim, n_update_dis, n_update_gen, 182 | update_max=update_max) 183 | 184 | return LapGan_model.generate(n_sample) 185 | 186 | 187 | if __name__ == '__main__': 188 | run_LAPGAN(n_epoch=1, D_featmap_dim=64, condi_D_featmap_dim=64, 189 | G_featmap_dim=128, condi_G_featmap_dim=64, update_max=50) 190 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Generative Adversarial Networks in PyTorch 2 | ======= 3 | 4 | 5 | [![Build Status](https://travis-ci.org/AaronYALai/Generative_Adversarial_Networks_PyTorch.svg?branch=master)](https://travis-ci.org/AaronYALai/Generative_Adversarial_Networks_PyTorch) 6 | [![Coverage Status](https://coveralls.io/repos/github/AaronYALai/Generative_Adversarial_Networks_PyTorch/badge.svg?branch=master)](https://coveralls.io/github/AaronYALai/Generative_Adversarial_Networks_PyTorch?branch=master) 7 | 8 | About 9 | -------- 10 | 11 | The repo is about the implementations of GAN, DCGAN, Improved GAN, LAPGAN, and InfoGAN in PyTorch. 12 | 13 | My presentation about GANs' recent development (at 2017.01.17): [Presentation slides](https://docs.google.com/presentation/d/1HRNjCo_0PlspynoJKuoEF1AYkaKaUNgMzQ4nqiTlNUM/edit#slide=id.p) 14 | 15 | Presented in the group meeting of Machine Discovery and Social Network Mining Lab, National Taiwan University. 16 | 17 | Content 18 | -------- 19 | 20 | - Generative Adversarial Nets (GAN): invented "adversarial nets" framework - a generative model G and a discriminative model D play a minimax two-player game. 21 | 22 | - DC-GAN: proposed a set of constraints on the architectural topology of Convolutional GANs that make them stable to train in most settings. 23 | 24 | - LAP-GAN: a cascade of generative models within a Laplacian pyramid framework to generate images in a coarse-to-fine fashion (high-resolution). 25 | 26 | - Improved GAN (minibatch discrimination): allow the discriminator to look at multiple data examples in combination by incorporating the closeness between examples in a minibatch as side information. 27 | 28 | - Info-GAN: an information-theoretic modification to the objective of Generative Adversarial Network that encourages it to learn interpretable and disentangled representations. 29 | 30 | Example 31 | --------- 32 | 33 | #### Sampled from MNIST dataset: 34 | 35 | 36 | #### Generated by DCGAN: 37 | 38 | 39 | #### Generated by Improved GAN: 40 | 41 | 42 | #### Generated by Info-GAN: 43 | 44 | 45 | Usage 46 | -------- 47 | Clone the repo and use the [virtualenv](http://www.virtualenv.org/): 48 | 49 | git clone https://github.com/AaronYALai/Generative_Adversarial_Networks_PyTorch 50 | 51 | cd Generative_Adversarial_Networks_PyTorch 52 | 53 | virtualenv venv 54 | 55 | source venv/bin/activate 56 | 57 | Install pytorch and all dependencies and run the model (in Linux): 58 | 59 | pip install https://download.pytorch.org/whl/cu75/torch-0.1.10.post2-cp27-none-linux_x86_64.whl 60 | 61 | pip install torchvision 62 | 63 | pip install -r requirements.txt 64 | 65 | cd GAN 66 | 67 | python run_GAN.py 68 | 69 | More details about the installation about PyTorch: 70 | 71 | 72 | References 73 | -------- 74 | 75 | - GAN: I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio, “Generative adversarial nets,” NIPS, 2014. 76 | 77 | - DC-GAN: Radford, Alec, Luke Metz, and Soumith Chintala. "Unsupervised representation learning with deep convolutional generative adversarial networks." arXiv 2015. 78 | 79 | - LAP-GAN: Denton, Emily L., Soumith Chintala, and Rob Fergus. "Deep Generative Image Models using a Laplacian Pyramid of Adversarial Networks." NIPS 2015. 80 | 81 | - Improved GAN: Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., & Chen, X. “Improved techniques for training gans.” NIPS 2016. 82 | 83 | - Info-GAN: Chen, X., Duan, Y., Houthooft, R., Schulman, J., Sutskever, I., & Abbeel, P. “Infogan: Interpretable representation learning by information maximizing generative adversarial nets.” NIPS 2016. 84 | -------------------------------------------------------------------------------- /images/DCGAN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AaronYALai/Generative_Adversarial_Networks_PyTorch/2e48957b650e9eea6bc941a8fe72035e69a32da2/images/DCGAN.png -------------------------------------------------------------------------------- /images/ImprovedGAN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AaronYALai/Generative_Adversarial_Networks_PyTorch/2e48957b650e9eea6bc941a8fe72035e69a32da2/images/ImprovedGAN.png -------------------------------------------------------------------------------- /images/InfoGAN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AaronYALai/Generative_Adversarial_Networks_PyTorch/2e48957b650e9eea6bc941a8fe72035e69a32da2/images/InfoGAN.png -------------------------------------------------------------------------------- /images/real.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AaronYALai/Generative_Adversarial_Networks_PyTorch/2e48957b650e9eea6bc941a8fe72035e69a32da2/images/real.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # Python wrapper of OpenCV: Open Source Computer Vision Library 2 | # https://pypi.python.org/pypi/opencv-python 3 | opencv-python==3.2.0.6 4 | 5 | # the modular source code checker: pep8, pyflakes and co 6 | # https://pypi.python.org/pypi/flake8/2.5.4 7 | flake8==2.5.4 8 | 9 | # pytest: simple powerful testing with Python 10 | # https://pypi.python.org/pypi/pytest/2.8.3 11 | pytest==2.8.3 12 | 13 | # Code coverage measurement for Python 14 | # https://pypi.python.org/pypi/coverage 15 | coverage==4.2 16 | 17 | # Pytest plugin for measuring coverage. 18 | # https://pypi.python.org/pypi/pytest-cov/2.2.0 19 | pytest-cov==2.2.0 20 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /tests/test_GANs.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | from GAN.run_GAN import run_GAN 3 | from DCGAN.run_DCGAN import run_DCGAN 4 | from ImprovedGAN.run_ImprovedGAN import run_ImprovedGAN 5 | from LAPGAN.run_LAPGAN import run_LAPGAN 6 | from InfoGAN.run_InfoGAN import run_InfoGAN 7 | 8 | 9 | class Test_running(TestCase): 10 | 11 | def test_run_GAN(self): 12 | run_GAN(n_epoch=1, batch_size=20, n_update_dis=5, n_update_gen=1, 13 | update_max=20) 14 | 15 | def test_run_DCGAN(self): 16 | run_DCGAN(n_epoch=1, batch_size=10, n_update_dis=1, n_update_gen=1, 17 | update_max=20) 18 | 19 | def test_run_ImprovedGAN(self): 20 | run_ImprovedGAN(n_epoch=1, batch_size=10, n_update_dis=1, 21 | n_update_gen=1, update_max=20) 22 | 23 | def test_run_LAPGAN(self): 24 | run_LAPGAN(n_epoch=1, batch_size=10, n_update_dis=1, 25 | n_update_gen=1, update_max=20) 26 | 27 | def test_run_InfoGAN(self): 28 | run_InfoGAN(n_conti=2, n_discrete=1, D_featmap_dim=64, 29 | G_featmap_dim=128, n_epoch=1, update_max=60) 30 | --------------------------------------------------------------------------------