├── img ├── CGAN.png ├── face1.png ├── face2.png ├── face3.png ├── face4.png ├── face5.png ├── DCGAN_36.png └── VAE-MNIST.gif ├── README.md ├── GANs ├── LSGAN-humanface.py ├── DCGAN-MNIST.py ├── WGAN-MNIST.py ├── LSGAN-MNIST.py ├── WGAN-GP-MNIST.py ├── WGAN-GP-humanface.py └── CGAN-MNIST.py └── VAE └── VAE-MNIST.py /img/CGAN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/CGAN.png -------------------------------------------------------------------------------- /img/face1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/face1.png -------------------------------------------------------------------------------- /img/face2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/face2.png -------------------------------------------------------------------------------- /img/face3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/face3.png -------------------------------------------------------------------------------- /img/face4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/face4.png -------------------------------------------------------------------------------- /img/face5.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/face5.png -------------------------------------------------------------------------------- /img/DCGAN_36.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/DCGAN_36.png -------------------------------------------------------------------------------- /img/VAE-MNIST.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/assassint2017/Generative-model-using-pytorch/HEAD/img/VAE-MNIST.gif -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Generative-model-GAN-VAE-using-pytorch 2 | 3 | implement **GANs** and **VAE** using **pytorch** 4 | 5 | GANs include **DCGAN LSGAN WGAN WGAN-GP** **CGAN**(use DCGAN architecture) 6 | VAE is from the paper **Auto-Encoding Variational Bayes** 7 | 8 | in my code, i use MNIST to test the network, but its very easy to switch dataset of you own 9 | 10 | here is result when i training DCGAN at 36 epoch 11 | 12 | ![image](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/DCGAN_36.png) 13 | 14 | second are some humanface generate by WGAN-GP 15 | 16 | the image use to train the WGAN is from baidu image 17 | 18 | ![face1](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/face1.png) 19 | ![face2](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/face2.png) 20 | ![face3](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/face3.png) 21 | ![face4](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/face4.png) 22 | ![face5](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/face5.png) 23 | 24 | next is the gif from training a VAE use MNIST datasset 25 | 26 | leftside is the reconstruction image,middle is the training image, and rightside is the image generate from the noise 27 | ![gif](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/VAE-MNIST.gif) 28 | 29 | next is the result from training a CDCGAN use MNIST dataset, each row represent a class of digit 30 | ![digit](https://github.com/assassint2017/Generative-model-GAN-VAE-using-pytorch/blob/master/img/CGAN.png) 31 | -------------------------------------------------------------------------------- /GANs/LSGAN-humanface.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import torch 4 | import torch.nn as nn 5 | from torch.utils.data import DataLoader 6 | from torch.autograd import Variable 7 | 8 | import torchvision 9 | import torchvision.transforms as transforms 10 | 11 | import cv2 as cv 12 | 13 | 14 | # Hyper parameters 15 | epoch_num = 200 16 | img_size = 64 # size of generated image 17 | batch_size = 128 18 | lr_g = 0.0002 # learning rate for Generator 19 | lr_d = 0.0002 # learning rate for Discriminator 20 | latent = 10 # dim of latent space 21 | img_channel = 3 # channel of generated image 22 | init_channel = 64 # control the initial Conv channel of the Generator and Discriminator 23 | workers = 1 # subprocess number for load the image 24 | k = 1 # train Discriminator K times and then train Generator one time 25 | 26 | mean = [0.5] 27 | std = [0.5] 28 | 29 | slope = 0.2 # slope for leaky relu 30 | 31 | # data enhancement 32 | data_transform = transforms.Compose([ 33 | transforms.Resize(size=img_size), 34 | transforms.RandomHorizontalFlip(), 35 | transforms.ToTensor(), 36 | transforms.Normalize(mean=mean, std=std) 37 | ]) 38 | 39 | # dataset 40 | dataset = torchvision.datasets.ImageFolder(root='./data/face/', transform=data_transform) 41 | 42 | data_loader = DataLoader(dataset, batch_size, True, num_workers=workers) 43 | 44 | 45 | # Generator 46 | class Generator(nn.Module): 47 | def __init__(self): 48 | super().__init__() 49 | 50 | self.deconv1 = nn.Sequential( 51 | nn.ConvTranspose2d(latent, init_channel * 8, 4, bias=False), 52 | nn.BatchNorm2d(init_channel * 8), 53 | nn.ReLU(), 54 | 55 | nn.Conv2d(init_channel * 8, init_channel * 8, 3, padding=1, bias=False), 56 | nn.BatchNorm2d(init_channel * 8), 57 | nn.ReLU() 58 | ) 59 | 60 | self.deconv2 = nn.Sequential( 61 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 62 | nn.BatchNorm2d(init_channel * 4), 63 | nn.ReLU(), 64 | 65 | nn.Conv2d(init_channel * 4, init_channel * 4, 3, padding=1, bias=False), 66 | nn.BatchNorm2d(init_channel * 4), 67 | nn.ReLU() 68 | ) 69 | 70 | self.deconv3 = nn.Sequential( 71 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 72 | nn.BatchNorm2d(init_channel * 2), 73 | nn.ReLU(), 74 | 75 | nn.Conv2d(init_channel * 2, init_channel * 2, 3, padding=1, bias=False), 76 | nn.BatchNorm2d(init_channel * 2), 77 | nn.ReLU() 78 | ) 79 | 80 | self.deconv4 = nn.Sequential( 81 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 82 | nn.BatchNorm2d(init_channel), 83 | nn.ReLU(), 84 | 85 | nn.Conv2d(init_channel, init_channel, 3, padding=1, bias=False), 86 | nn.BatchNorm2d(init_channel), 87 | nn.ReLU() 88 | ) 89 | 90 | self.deconv5 = nn.Sequential( 91 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 92 | nn.Tanh() 93 | ) 94 | 95 | # initialization for parameters 96 | 97 | for layer in self.modules(): 98 | 99 | if isinstance(layer, nn.ConvTranspose2d): 100 | nn.init.normal(layer.weight.data, 0, 0.02) 101 | 102 | elif isinstance(layer, nn.BatchNorm2d): 103 | layer.weight.data.fill_(1) 104 | layer.bias.data.zero_() 105 | 106 | def forward(self, inputs): 107 | 108 | outputs = self.deconv1(inputs) 109 | outputs = self.deconv2(outputs) 110 | outputs = self.deconv3(outputs) 111 | outputs = self.deconv4(outputs) 112 | outputs = self.deconv5(outputs) 113 | 114 | return outputs 115 | 116 | 117 | # Discriminator( 118 | class Discriminator(nn.Module): 119 | def __init__(self): 120 | super().__init__() 121 | 122 | self.conv1 = nn.Sequential( 123 | nn.Conv2d(img_channel, init_channel, 4, 2, 1, bias=False), 124 | nn.LeakyReLU(slope) 125 | ) 126 | 127 | self.conv2 = nn.Sequential( 128 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 129 | nn.BatchNorm2d(init_channel * 2), 130 | nn.LeakyReLU(slope) 131 | ) 132 | 133 | self.conv3 = nn.Sequential( 134 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 135 | nn.BatchNorm2d(init_channel * 4), 136 | nn.LeakyReLU(slope) 137 | ) 138 | 139 | self.conv4 = nn.Sequential( 140 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 141 | nn.BatchNorm2d(init_channel * 8), 142 | nn.LeakyReLU(slope) 143 | ) 144 | 145 | self.conv5 = nn.Sequential( 146 | 147 | # no sigmoid! 148 | nn.Conv2d(init_channel * 8, 1, 4, bias=False) 149 | ) 150 | 151 | # initialization for parameters 152 | for layer in self.modules(): 153 | 154 | if isinstance(layer, nn.Conv2d): 155 | nn.init.normal(layer.weight.data, 0, 0.02) 156 | 157 | elif isinstance(layer, nn.BatchNorm2d): 158 | layer.weight.data.fill_(1) 159 | layer.bias.data.zero_() 160 | 161 | def forward(self, inputs): 162 | 163 | outputs = self.conv1(inputs) 164 | outputs = self.conv2(outputs) 165 | outputs = self.conv3(outputs) 166 | outputs = self.conv4(outputs) 167 | outputs = self.conv5(outputs) 168 | 169 | return outputs.view(inputs.size(0)) 170 | 171 | 172 | # use cuda if you have GPU 173 | net_g = Generator().cuda() 174 | net_d = Discriminator().cuda() 175 | 176 | # optimizer 177 | opt_g = torch.optim.Adam(net_g.parameters(), lr=lr_g, betas=(0.5, 0.999)) # optimizer for Generator 178 | opt_d = torch.optim.Adam(net_d.parameters(), lr=lr_d, betas=(0.5, 0.999)) # optimizer for Discriminator 179 | 180 | 181 | # get random noise 182 | def get_noise(noise_num=batch_size): 183 | 184 | return Variable(torch.randn((noise_num, latent, 1, 1)).cuda()) 185 | 186 | 187 | # train the network 188 | start = time() 189 | fix_noise = get_noise(64) 190 | number = 1 191 | 192 | for epoch in range(epoch_num): 193 | 194 | for step, (real_data, target) in enumerate(data_loader, 1): 195 | 196 | # train Discriminator 197 | 198 | real_data = Variable(real_data).cuda() 199 | 200 | prob_fake = net_d(net_g(get_noise(real_data.size(0)))) 201 | 202 | prob_real = net_d(real_data) 203 | 204 | loss_d = 0.5 * torch.mean(torch.pow(prob_real - 1, 2) + torch.pow(prob_fake, 2)) 205 | 206 | opt_d.zero_grad() 207 | loss_d.backward() 208 | opt_d.step() 209 | 210 | # train Generator 211 | if step % k is 0: 212 | 213 | prob_fake = net_d(net_g(get_noise())) 214 | 215 | loss_g = 0.5 * torch.mean(torch.pow(prob_fake - 1, 2)) 216 | 217 | opt_g.zero_grad() 218 | loss_g.backward() 219 | opt_g.step() 220 | 221 | if step % 20 is 0: 222 | 223 | fake_img = torchvision.utils.make_grid((0.5 * net_g(fix_noise).data.cpu() + 0.5)) 224 | 225 | cv.imwrite('./img/LS/' + str(number) + '.png', fake_img.permute(1, 2, 0).numpy() * 255) 226 | 227 | number += 1 228 | -------------------------------------------------------------------------------- /GANs/DCGAN-MNIST.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import math 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.utils.data import DataLoader 8 | from torch.autograd import Variable 9 | 10 | import torchvision 11 | import torchvision.transforms as transforms 12 | 13 | import matplotlib.pyplot as plt 14 | 15 | import tensorboardX 16 | 17 | # Hyper parameters 18 | epoch_num = 200 19 | img_size = 64 # size of generated image 20 | batch_size = 128 21 | lr_g = 0.0002 # learning rate for Generator 22 | lr_d = 0.0002 # learning rate for Discriminator 23 | latent = 100 # dim of latent space 24 | img_channel = 1 # channel of generated image 25 | init_channel = 16 # control the initial Conv channel of the Generator and Discriminator 26 | workers = 2 # subprocess number for load the image 27 | k = 1 # train Discriminator K times and then train Generator one time 28 | dataset_size = 60000 # image number of your training set 29 | 30 | mean = [0.5] 31 | std = [0.5] 32 | 33 | slope = 0.2 # slope for leaky relu 34 | 35 | # use tensorboard 36 | writer = tensorboardX.SummaryWriter(log_dir='./logs/') 37 | 38 | # data enhancement 39 | data_transform = transforms.Compose([ 40 | transforms.Resize(size=img_size), 41 | transforms.RandomRotation(5), 42 | transforms.RandomHorizontalFlip(), 43 | transforms.ToTensor(), 44 | transforms.Normalize(mean=mean, std=std) 45 | ]) 46 | 47 | # dataset 48 | data_set = torchvision.datasets.MNIST( 49 | root='./MNIST', 50 | train=True, 51 | transform=data_transform 52 | ) 53 | 54 | data_loader = DataLoader(data_set, batch_size, True, num_workers=workers) 55 | 56 | 57 | # Generator 58 | class Generator(nn.Module): 59 | def __init__(self): 60 | super().__init__() 61 | 62 | self.deconv1 = nn.Sequential( 63 | nn.ConvTranspose2d(latent, init_channel * 8, 4, bias=False), 64 | nn.BatchNorm2d(init_channel * 8), 65 | nn.ReLU() 66 | ) 67 | 68 | self.deconv2 = nn.Sequential( 69 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 70 | nn.BatchNorm2d(init_channel * 4), 71 | nn.ReLU() 72 | ) 73 | 74 | self.deconv3 = nn.Sequential( 75 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 76 | nn.BatchNorm2d(init_channel * 2), 77 | nn.ReLU() 78 | ) 79 | 80 | self.deconv4 = nn.Sequential( 81 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 82 | nn.BatchNorm2d(init_channel), 83 | nn.ReLU() 84 | ) 85 | 86 | self.deconv5 = nn.Sequential( 87 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 88 | nn.Tanh() 89 | ) 90 | 91 | # initialization for parameters 92 | 93 | for layer in self.modules(): 94 | 95 | if isinstance(layer, nn.ConvTranspose2d): 96 | nn.init.normal(layer.weight.data, 0, 0.02) 97 | 98 | elif isinstance(layer, nn.BatchNorm2d): 99 | layer.weight.data.fill_(1) 100 | layer.bias.data.zero_() 101 | 102 | def forward(self, inputs): 103 | 104 | outputs = self.deconv1(inputs) 105 | outputs = self.deconv2(outputs) 106 | outputs = self.deconv3(outputs) 107 | outputs = self.deconv4(outputs) 108 | outputs = self.deconv5(outputs) 109 | 110 | return outputs 111 | 112 | 113 | # Discriminator( 114 | class Discriminator(nn.Module): 115 | def __init__(self): 116 | super().__init__() 117 | 118 | self.conv1 = nn.Sequential( 119 | nn.Conv2d(img_channel, init_channel, 4, 2, 1, bias=False), 120 | nn.LeakyReLU(slope) 121 | ) 122 | 123 | self.conv2 = nn.Sequential( 124 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 125 | nn.BatchNorm2d(init_channel * 2), 126 | nn.LeakyReLU(slope) 127 | ) 128 | 129 | self.conv3 = nn.Sequential( 130 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 131 | nn.BatchNorm2d(init_channel * 4), 132 | nn.LeakyReLU(slope) 133 | ) 134 | 135 | self.conv4 = nn.Sequential( 136 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 137 | nn.BatchNorm2d(init_channel * 8), 138 | nn.LeakyReLU(slope) 139 | ) 140 | 141 | self.conv5 = nn.Sequential( 142 | nn.Conv2d(init_channel * 8, 1, 4, bias=False), 143 | nn.Sigmoid() 144 | ) 145 | 146 | # initialization for parameters 147 | for layer in self.modules(): 148 | 149 | if isinstance(layer, nn.Conv2d): 150 | nn.init.normal(layer.weight.data, 0, 0.02) 151 | 152 | elif isinstance(layer, nn.BatchNorm2d): 153 | layer.weight.data.fill_(1) 154 | layer.bias.data.zero_() 155 | 156 | def forward(self, inputs): 157 | 158 | outputs = self.conv1(inputs) 159 | outputs = self.conv2(outputs) 160 | outputs = self.conv3(outputs) 161 | outputs = self.conv4(outputs) 162 | outputs = self.conv5(outputs) 163 | 164 | return outputs.view(inputs.size(0)) 165 | 166 | 167 | # use cuda if you have GPU 168 | net_g = Generator().cuda() 169 | net_d = Discriminator().cuda() 170 | 171 | # use tensorboard draw the computational graph 172 | writer.add_graph(net_g, Variable(torch.randn(batch_size, latent, 1, 1).cuda())) 173 | 174 | writer.add_graph(net_d, Variable(torch.FloatTensor(batch_size, img_channel, img_size, img_size).cuda())) 175 | 176 | # optimizer 177 | opt_g = torch.optim.Adam(net_g.parameters(), lr=lr_g, betas=(0.5, 0.999)) # optimizer for Generator 178 | opt_d = torch.optim.Adam(net_d.parameters(), lr=lr_d, betas=(0.5, 0.999)) # optimizer for Discriminator 179 | 180 | 181 | # get random noise 182 | def get_noise(noise_num=batch_size): 183 | 184 | return Variable(torch.randn((noise_num, latent, 1, 1)).cuda()) 185 | 186 | 187 | # train the network 188 | start = time() 189 | fix_noise = get_noise(64) 190 | 191 | for epoch in range(epoch_num): 192 | 193 | for step, (real_data, target) in enumerate(data_loader, 1): 194 | 195 | # train Discriminator 196 | 197 | real_data = Variable(real_data).cuda() 198 | 199 | prob_fake = net_d(net_g(get_noise(real_data.size(0)))) 200 | 201 | prob_real = net_d(real_data) 202 | 203 | loss_d = - torch.mean(torch.log(prob_real) + torch.log(1 - prob_fake)) 204 | 205 | opt_d.zero_grad() 206 | loss_d.backward() 207 | opt_d.step() 208 | 209 | # train Generator 210 | if step % k is 0: 211 | 212 | prob_fake = net_d(net_g(get_noise())) 213 | 214 | loss_g = - torch.mean(torch.log(prob_fake)) 215 | 216 | opt_g.zero_grad() 217 | loss_g.backward() 218 | opt_g.step() 219 | 220 | if step % 20 is 0: 221 | 222 | iteration = epoch * (math.ceil(dataset_size / batch_size)) + step 223 | writer.add_scalars('loss', {'g_loss': loss_g.data[0], 'd_loss': loss_d.data[0]}, iteration) 224 | print('epoch:', epoch, 'step', step, 'time:', (time() - start) / 60, 'min') 225 | 226 | fake_img = torchvision.utils.make_grid((0.5 * net_g(fix_noise).data.cpu() + 0.5)) 227 | plt.imshow(fake_img.permute(1, 2, 0).numpy()) 228 | plt.pause(0.01) 229 | 230 | plt.show() 231 | -------------------------------------------------------------------------------- /GANs/WGAN-MNIST.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import math 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.utils.data import DataLoader 8 | from torch.autograd import Variable 9 | 10 | import torchvision 11 | import torchvision.transforms as transforms 12 | 13 | import matplotlib.pyplot as plt 14 | 15 | import tensorboardX 16 | 17 | from my_dataset import MyDataSet 18 | 19 | # Hyper parameters 20 | epoch_num = 200 21 | img_size = 64 # size of generated image 22 | batch_size = 128 23 | lr_g = 0.0002 # learning rate for Generator 24 | lr_c = 0.0002 # learning rate for Critic 25 | latent = 100 # dim of latent space 26 | img_channel = 1 # channel of generated image 27 | init_channel = 16 # control the initial Conv channel of the Generator and Discriminator 28 | workers = 2 # subprocess number for load the image 29 | k = 5 # train Critic K times and then train Generator one time 30 | dataset_size = 60000 # image number of your training set 31 | 32 | mean = [0.5] 33 | std = [0.5] 34 | 35 | slope = 0.2 # slope for leaky relu 36 | 37 | c = 0.01 # weight clipping 38 | 39 | # use tensorboard 40 | writer = tensorboardX.SummaryWriter(log_dir='./logs/') 41 | 42 | # data enhancement 43 | data_transform = transforms.Compose([ 44 | transforms.Resize(size=img_size), 45 | transforms.RandomRotation(5), 46 | transforms.RandomHorizontalFlip(), 47 | transforms.ToTensor(), 48 | transforms.Normalize(mean=mean, std=std) 49 | ]) 50 | 51 | # dataset 52 | data_set = torchvision.datasets.MNIST( 53 | root='./MNIST', 54 | train=True, 55 | transform=data_transform 56 | ) 57 | 58 | data_loader = DataLoader(data_set, batch_size, True, num_workers=workers) 59 | 60 | 61 | # Generator 62 | class Generator(nn.Module): 63 | def __init__(self): 64 | super().__init__() 65 | 66 | self.deconv1 = nn.Sequential( 67 | nn.ConvTranspose2d(latent, init_channel * 8, 4, bias=False), 68 | nn.BatchNorm2d(init_channel * 8), 69 | nn.ReLU() 70 | ) 71 | 72 | self.deconv2 = nn.Sequential( 73 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 74 | nn.BatchNorm2d(init_channel * 4), 75 | nn.ReLU() 76 | ) 77 | 78 | self.deconv3 = nn.Sequential( 79 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 80 | nn.BatchNorm2d(init_channel * 2), 81 | nn.ReLU() 82 | ) 83 | 84 | self.deconv4 = nn.Sequential( 85 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 86 | nn.BatchNorm2d(init_channel), 87 | nn.ReLU() 88 | ) 89 | 90 | self.deconv5 = nn.Sequential( 91 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 92 | nn.Tanh() 93 | ) 94 | 95 | # initialization for parameters 96 | 97 | for layer in self.modules(): 98 | 99 | if isinstance(layer, nn.ConvTranspose2d): 100 | 101 | nn.init.normal(layer.weight.data, 0, 0.02) 102 | 103 | elif isinstance(layer, nn.BatchNorm2d): 104 | layer.weight.data.fill_(1) 105 | layer.bias.data.zero_() 106 | 107 | def forward(self, inputs): 108 | 109 | outputs = self.deconv1(inputs) 110 | outputs = self.deconv2(outputs) 111 | outputs = self.deconv3(outputs) 112 | outputs = self.deconv4(outputs) 113 | outputs = self.deconv5(outputs) 114 | 115 | return outputs 116 | 117 | 118 | # Critic( 119 | class Critic(nn.Module): 120 | def __init__(self): 121 | super().__init__() 122 | 123 | self.conv1 = nn.Sequential( 124 | nn.Conv2d(img_channel, init_channel, 4, 2, 1, bias=False), 125 | nn.LeakyReLU(slope) 126 | ) 127 | 128 | self.conv2 = nn.Sequential( 129 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 130 | nn.BatchNorm2d(init_channel * 2), 131 | nn.LeakyReLU(slope) 132 | ) 133 | 134 | self.conv3 = nn.Sequential( 135 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 136 | nn.BatchNorm2d(init_channel * 4), 137 | nn.LeakyReLU(slope) 138 | ) 139 | 140 | self.conv4 = nn.Sequential( 141 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 142 | nn.BatchNorm2d(init_channel * 8), 143 | nn.LeakyReLU(slope) 144 | ) 145 | 146 | self.conv5 = nn.Sequential( 147 | 148 | # no sigmoid! 149 | nn.Conv2d(init_channel * 8, 1, 4, bias=False) 150 | ) 151 | 152 | # initialization for parameters 153 | for layer in self.modules(): 154 | 155 | if isinstance(layer, nn.Conv2d): 156 | nn.init.normal(layer.weight.data, 0, 0.02) 157 | 158 | elif isinstance(layer, nn.BatchNorm2d): 159 | layer.weight.data.fill_(1) 160 | layer.bias.data.zero_() 161 | 162 | def forward(self, inputs): 163 | 164 | outputs = self.conv1(inputs) 165 | outputs = self.conv2(outputs) 166 | outputs = self.conv3(outputs) 167 | outputs = self.conv4(outputs) 168 | outputs = self.conv5(outputs) 169 | 170 | return outputs.view(inputs.size(0)) 171 | 172 | 173 | # use cuda if you have GPU 174 | net_g = Generator().cuda() 175 | net_c = Critic().cuda() 176 | 177 | 178 | # use tensorboard draw the computational graph 179 | writer.add_graph(net_g, Variable(torch.randn(batch_size, latent, 1, 1).cuda())) 180 | 181 | writer.add_graph(net_c, Variable(torch.FloatTensor(batch_size, img_channel, img_size, img_size).cuda())) 182 | 183 | # optimizer 184 | opt_g = torch.optim.RMSprop(net_g.parameters(), lr=lr_g) # optimizer for Generator 185 | opt_c = torch.optim.RMSprop(net_c.parameters(), lr=lr_c) # optimizer for Critic 186 | 187 | 188 | # get random noise 189 | def get_noise(noise_num=batch_size): 190 | 191 | return Variable(torch.randn((noise_num, latent, 1, 1)).cuda()) 192 | 193 | 194 | # train the network 195 | start = time() 196 | fix_noise = get_noise(64) 197 | 198 | for epoch in range(epoch_num): 199 | 200 | for step, (real_data, target) in enumerate(data_loader, 1): 201 | 202 | # weight clipping 203 | for parm in net_c.parameters(): 204 | parm.data.clamp_(-c, c) 205 | 206 | # train Critic 207 | 208 | real_data = Variable(real_data).cuda() 209 | 210 | output_fake = net_c(net_g(get_noise(real_data.size(0)))) 211 | 212 | output_real = net_c(real_data) 213 | 214 | loss_c = torch.mean(output_fake - output_real) 215 | 216 | opt_c.zero_grad() 217 | loss_c.backward() 218 | opt_c.step() 219 | 220 | # train Generator 221 | if step % k is 0: 222 | 223 | output_fake = net_c(net_g(get_noise())) 224 | 225 | loss_g = - torch.mean(output_fake) 226 | 227 | opt_g.zero_grad() 228 | loss_g.backward() 229 | opt_g.step() 230 | 231 | if step % 20 is 0: 232 | 233 | iteration = epoch * (math.ceil(dataset_size / batch_size)) + step 234 | writer.add_scalars('loss', {'g_loss': loss_g.data[0], 'c_loss': loss_c.data[0]}, iteration) 235 | print('epoch:', epoch, 'step', step, 'time:', (time() - start) / 60, 'min') 236 | 237 | fake_img = torchvision.utils.make_grid((0.5 * net_g(fix_noise).data.cpu() + 0.5)) 238 | plt.imshow(fake_img.permute(1, 2, 0).numpy()) 239 | plt.pause(0.01) 240 | 241 | plt.show() 242 | 243 | -------------------------------------------------------------------------------- /GANs/LSGAN-MNIST.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import math 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.utils.data import DataLoader 8 | from torch.autograd import Variable 9 | 10 | import torchvision 11 | import torchvision.transforms as transforms 12 | 13 | import matplotlib.pyplot as plt 14 | 15 | import tensorboardX 16 | 17 | 18 | # Hyper parameters 19 | epoch_num = 200 20 | img_size = 64 # size of generated image 21 | batch_size = 128 22 | lr_g = 0.0002 # learning rate for Generator 23 | lr_d = 0.0002 # learning rate for Discriminator 24 | latent = 10 # dim of latent space 25 | img_channel = 3 # channel of generated image 26 | init_channel = 64 # control the initial Conv channel of the Generator and Discriminator 27 | workers = 2 # subprocess number for load the image 28 | k = 1 # train Discriminator K times and then train Generator one time 29 | dataset_size = 50000 # image number of your training set 30 | 31 | mean = [0.5] 32 | std = [0.5] 33 | 34 | slope = 0.2 # slope for leaky relu 35 | 36 | # use tensorboard 37 | writer = tensorboardX.SummaryWriter(log_dir='./logs/') 38 | 39 | # data enhancement 40 | data_transform = transforms.Compose([ 41 | transforms.Resize(size=img_size), 42 | transforms.RandomRotation(5), 43 | transforms.RandomHorizontalFlip(), 44 | transforms.ToTensor(), 45 | transforms.Normalize(mean=mean, std=std) 46 | ]) 47 | 48 | # dataset 49 | data_set = torchvision.datasets.CIFAR10( 50 | root='./CIFAR10', 51 | train=True, 52 | transform=data_transform 53 | ) 54 | 55 | data_loader = DataLoader(data_set, batch_size, True, num_workers=workers) 56 | 57 | 58 | # Generator 59 | class Generator(nn.Module): 60 | def __init__(self): 61 | super().__init__() 62 | 63 | self.deconv1 = nn.Sequential( 64 | nn.ConvTranspose2d(latent, init_channel * 8, 4, bias=False), 65 | nn.BatchNorm2d(init_channel * 8), 66 | nn.ReLU() 67 | ) 68 | 69 | self.deconv2 = nn.Sequential( 70 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 71 | nn.BatchNorm2d(init_channel * 4), 72 | nn.ReLU() 73 | ) 74 | 75 | self.deconv3 = nn.Sequential( 76 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 77 | nn.BatchNorm2d(init_channel * 2), 78 | nn.ReLU() 79 | ) 80 | 81 | self.deconv4 = nn.Sequential( 82 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 83 | nn.BatchNorm2d(init_channel), 84 | nn.ReLU() 85 | ) 86 | 87 | self.deconv5 = nn.Sequential( 88 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 89 | nn.Tanh() 90 | ) 91 | 92 | # initialization for parameters 93 | 94 | for layer in self.modules(): 95 | 96 | if isinstance(layer, nn.ConvTranspose2d): 97 | nn.init.normal(layer.weight.data, 0, 0.02) 98 | 99 | elif isinstance(layer, nn.BatchNorm2d): 100 | layer.weight.data.fill_(1) 101 | layer.bias.data.zero_() 102 | 103 | def forward(self, inputs): 104 | 105 | outputs = self.deconv1(inputs) 106 | outputs = self.deconv2(outputs) 107 | outputs = self.deconv3(outputs) 108 | outputs = self.deconv4(outputs) 109 | outputs = self.deconv5(outputs) 110 | 111 | return outputs 112 | 113 | 114 | # Discriminator( 115 | class Discriminator(nn.Module): 116 | def __init__(self): 117 | super().__init__() 118 | 119 | self.conv1 = nn.Sequential( 120 | nn.Conv2d(img_channel, init_channel, 4, 2, 1, bias=False), 121 | nn.LeakyReLU(slope) 122 | ) 123 | 124 | self.conv2 = nn.Sequential( 125 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 126 | nn.BatchNorm2d(init_channel * 2), 127 | nn.LeakyReLU(slope) 128 | ) 129 | 130 | self.conv3 = nn.Sequential( 131 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 132 | nn.BatchNorm2d(init_channel * 4), 133 | nn.LeakyReLU(slope) 134 | ) 135 | 136 | self.conv4 = nn.Sequential( 137 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 138 | nn.BatchNorm2d(init_channel * 8), 139 | nn.LeakyReLU(slope) 140 | ) 141 | 142 | self.conv5 = nn.Sequential( 143 | 144 | # no sigmoid! 145 | nn.Conv2d(init_channel * 8, 1, 4, bias=False) 146 | ) 147 | 148 | # initialization for parameters 149 | for layer in self.modules(): 150 | 151 | if isinstance(layer, nn.Conv2d): 152 | nn.init.normal(layer.weight.data, 0, 0.02) 153 | 154 | elif isinstance(layer, nn.BatchNorm2d): 155 | layer.weight.data.fill_(1) 156 | layer.bias.data.zero_() 157 | 158 | def forward(self, inputs): 159 | 160 | outputs = self.conv1(inputs) 161 | outputs = self.conv2(outputs) 162 | outputs = self.conv3(outputs) 163 | outputs = self.conv4(outputs) 164 | outputs = self.conv5(outputs) 165 | 166 | return outputs.view(inputs.size(0)) 167 | 168 | 169 | # use cuda if you have GPU 170 | net_g = Generator().cuda() 171 | net_d = Discriminator().cuda() 172 | 173 | # use tensorboard draw the computational graph 174 | writer.add_graph(net_g, Variable(torch.randn(batch_size, latent, 1, 1).cuda())) 175 | 176 | writer.add_graph(net_d, Variable(torch.FloatTensor(batch_size, img_channel, img_size, img_size).cuda())) 177 | 178 | # optimizer 179 | opt_g = torch.optim.Adam(net_g.parameters(), lr=lr_g, betas=(0.5, 0.999)) # optimizer for Generator 180 | opt_d = torch.optim.Adam(net_d.parameters(), lr=lr_d, betas=(0.5, 0.999)) # optimizer for Discriminator 181 | 182 | 183 | # get random noise 184 | def get_noise(noise_num=batch_size): 185 | 186 | return Variable(torch.randn((noise_num, latent, 1, 1)).cuda()) 187 | 188 | 189 | # train the network 190 | start = time() 191 | fix_noise = get_noise(64) 192 | number = 1 193 | 194 | for epoch in range(epoch_num): 195 | 196 | for step, (real_data, target) in enumerate(data_loader, 1): 197 | 198 | # train Discriminator 199 | 200 | real_data = Variable(real_data).cuda() 201 | 202 | prob_fake = net_d(net_g(get_noise(real_data.size(0)))) 203 | 204 | prob_real = net_d(real_data) 205 | 206 | loss_d = 0.5 * torch.mean(torch.pow(prob_real - 1, 2) + torch.pow(prob_fake, 2)) 207 | 208 | opt_d.zero_grad() 209 | loss_d.backward() 210 | opt_d.step() 211 | 212 | # train Generator 213 | if step % k is 0: 214 | 215 | prob_fake = net_d(net_g(get_noise())) 216 | 217 | loss_g = 0.5 * torch.mean(torch.pow(prob_fake - 1, 2)) 218 | 219 | opt_g.zero_grad() 220 | loss_g.backward() 221 | opt_g.step() 222 | 223 | if step % 20 is 0: 224 | 225 | iteration = epoch * (math.ceil(dataset_size / batch_size)) + step 226 | writer.add_scalars('loss', {'g_loss': loss_g.data[0], 'd_loss': loss_d.data[0]}, iteration) 227 | print('epoch:', epoch, 'step', step, 'time:', (time() - start) / 60, 'min') 228 | 229 | fake_img = torchvision.utils.make_grid((0.5 * net_g(fix_noise).data.cpu() + 0.5)) 230 | plt.imshow(fake_img.permute(1, 2, 0).numpy()) 231 | plt.pause(0.01) 232 | 233 | #plt.imsave('./img/'+str(number)+'.png', fake_img.permute(1, 2, 0).numpy()) 234 | number += 1 235 | 236 | plt.show() 237 | -------------------------------------------------------------------------------- /VAE/VAE-MNIST.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import math 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.utils.data import DataLoader 8 | from torch.autograd import Variable 9 | 10 | import torchvision 11 | import torchvision.transforms as transforms 12 | 13 | import matplotlib.pyplot as plt 14 | 15 | 16 | # Hyper parameters 17 | epoch_num = 200 18 | img_size = 64 # size of generated image 19 | batch_size = 128 20 | lr_en = 0.0002 # learning rate for encoder 21 | lr_de = 0.0002 # learning rate for decoder 22 | latent = 100 # dim of latent space 23 | img_channel = 1 # channel of generated image 24 | init_channel = 16 # control the initial Conv channel of the Generator and Discriminator 25 | workers = 1 # subprocess number for load the image 26 | dataset_size = 60000 # image number of your training set 27 | 28 | mean = [0.5] 29 | std = [0.5] 30 | 31 | slope = 0.2 # slope for leaky relu 32 | 33 | # data enhancement 34 | data_transform = transforms.Compose([ 35 | transforms.Resize(size=img_size), 36 | transforms.ToTensor(), 37 | transforms.Normalize(mean=mean, std=std) 38 | ]) 39 | 40 | # dataset 41 | data_set = torchvision.datasets.MNIST( 42 | root='./MNIST', 43 | train=True, 44 | transform=data_transform 45 | ) 46 | 47 | data_loader = DataLoader(data_set, batch_size, True, num_workers=workers) 48 | 49 | 50 | # Encoder 51 | class Encoder(nn.Module): 52 | def __init__(self): 53 | super().__init__() 54 | 55 | self.conv1 = nn.Sequential( 56 | nn.Conv2d(img_channel, init_channel, 4, 2, 1, bias=False), 57 | nn.LeakyReLU(slope) 58 | ) 59 | 60 | self.conv2 = nn.Sequential( 61 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 62 | nn.BatchNorm2d(init_channel * 2), 63 | nn.LeakyReLU(slope) 64 | ) 65 | 66 | self.conv3 = nn.Sequential( 67 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 68 | nn.BatchNorm2d(init_channel * 4), 69 | nn.LeakyReLU(slope) 70 | ) 71 | 72 | self.conv4 = nn.Sequential( 73 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 74 | nn.BatchNorm2d(init_channel * 8), 75 | nn.LeakyReLU(slope) 76 | ) 77 | 78 | self.conv5 = nn.Conv2d(init_channel * 8, init_channel * 8, 4, bias=False) 79 | 80 | self.mean = nn.Linear(init_channel * 8, latent) 81 | 82 | self.log_var = nn.Linear(init_channel * 8, latent) 83 | 84 | # initialization for parameters 85 | for layer in self.modules(): 86 | 87 | if isinstance(layer, nn.Conv2d): 88 | nn.init.normal(layer.weight.data, 0, 0.02) 89 | 90 | elif isinstance(layer, nn.BatchNorm2d): 91 | layer.weight.data.fill_(1) 92 | layer.bias.data.zero_() 93 | 94 | def forward(self, inputs): 95 | 96 | outputs = self.conv1(inputs) 97 | outputs = self.conv2(outputs) 98 | outputs = self.conv3(outputs) 99 | outputs = self.conv4(outputs) 100 | outputs = self.conv5(outputs) 101 | 102 | outputs = outputs.view(inputs.size(0), -1) 103 | 104 | return self.mean(outputs), self.log_var(outputs) 105 | 106 | 107 | # Decoder 108 | class Decoder(nn.Module): 109 | def __init__(self): 110 | super().__init__() 111 | 112 | self.deconv1 = nn.Sequential( 113 | nn.ConvTranspose2d(latent, init_channel * 8, 4, bias=False), 114 | nn.BatchNorm2d(init_channel * 8), 115 | nn.ReLU(), 116 | ) 117 | 118 | self.deconv2 = nn.Sequential( 119 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 120 | nn.BatchNorm2d(init_channel * 4), 121 | nn.ReLU(), 122 | ) 123 | 124 | self.deconv3 = nn.Sequential( 125 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 126 | nn.BatchNorm2d(init_channel * 2), 127 | nn.ReLU(), 128 | ) 129 | 130 | self.deconv4 = nn.Sequential( 131 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 132 | nn.BatchNorm2d(init_channel), 133 | nn.ReLU(), 134 | ) 135 | 136 | self.deconv5 = nn.Sequential( 137 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 138 | nn.Tanh() 139 | ) 140 | 141 | # initialization for parameters 142 | 143 | for layer in self.modules(): 144 | 145 | if isinstance(layer, nn.ConvTranspose2d): 146 | nn.init.normal(layer.weight.data, 0, 0.02) 147 | 148 | elif isinstance(layer, nn.BatchNorm2d): 149 | layer.weight.data.fill_(1) 150 | layer.bias.data.zero_() 151 | 152 | def forward(self, inputs): 153 | 154 | inputs = inputs.view(inputs.size(0), inputs.size(1), 1, 1) 155 | 156 | outputs = self.deconv1(inputs) 157 | outputs = self.deconv2(outputs) 158 | outputs = self.deconv3(outputs) 159 | outputs = self.deconv4(outputs) 160 | outputs = self.deconv5(outputs) 161 | 162 | return outputs 163 | 164 | 165 | # use cuda if you have GPU 166 | decoder = Decoder().cuda() 167 | encoder = Encoder().cuda() 168 | 169 | # optimizer 170 | opt_d = torch.optim.Adam(decoder.parameters(), lr=lr_de) # optimizer for decoder 171 | opt_e = torch.optim.Adam(encoder.parameters(), lr=lr_en) # optimizer for encoder 172 | 173 | 174 | # get random noise 175 | def get_noise(noise_num=batch_size): 176 | 177 | return Variable(torch.randn((noise_num, latent)).cuda()) 178 | 179 | 180 | # reconstruction loss 181 | re_loss = nn.MSELoss() 182 | 183 | # train the network 184 | start = time() 185 | number = 1 186 | img = plt.figure('Visualization') 187 | 188 | for epoch in range(epoch_num): 189 | 190 | for step, (real_data, target) in enumerate(data_loader, 1): 191 | 192 | opt_e.zero_grad() 193 | opt_d.zero_grad() 194 | 195 | real_data = Variable(real_data).cuda() 196 | 197 | mean, log_var = encoder(real_data) 198 | 199 | fake_noise = mean + (torch.exp(log_var / 2) * get_noise(real_data.size(0))) 200 | 201 | reconstruct_img = decoder(fake_noise) 202 | 203 | KL_loss = torch.sum(- 0.5 * (1 + log_var ** 2 - mean ** 2 - torch.exp(log_var)), dim=1) 204 | 205 | loss = re_loss(reconstruct_img, real_data) + torch.mean(KL_loss) 206 | 207 | loss.backward() 208 | 209 | opt_d.step() 210 | opt_e.step() 211 | 212 | if step % 20 is 0: 213 | 214 | iteration = epoch * (math.ceil(dataset_size / batch_size)) + step 215 | print('epoch:', epoch, 'step', step, 'time:', (time() - start) / 60, 'min') 216 | 217 | img.add_subplot(131) 218 | plt.title('reconstruction_img') 219 | 220 | re_img = torchvision.utils.make_grid(reconstruct_img.data[0:36].cpu() * 0.5 + 0.5, nrow=6) 221 | plt.imshow(re_img.permute(1, 2, 0).numpy()) 222 | 223 | img.add_subplot(132) 224 | plt.title('train_img') 225 | 226 | real_img = torchvision.utils.make_grid(real_data.data[0:36].cpu() * 0.5 + 0.5, nrow=6) 227 | plt.imshow(real_img.permute(1, 2, 0).numpy()) 228 | 229 | img.add_subplot(133) 230 | plt.title('generate_img') 231 | 232 | ge_img = torchvision.utils.make_grid(decoder(get_noise(36)).data.cpu() * 0.5 + 0.5, nrow=6) 233 | plt.imshow(ge_img.permute(1, 2, 0).numpy()) 234 | 235 | img.savefig('./img/' + str(number) + '.png') 236 | 237 | number += 1 238 | 239 | plt.pause(0.01) 240 | 241 | -------------------------------------------------------------------------------- /GANs/WGAN-GP-MNIST.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import math 4 | 5 | import torch 6 | import torch.nn as nn 7 | from torch.utils.data import DataLoader 8 | from torch.autograd import Variable 9 | 10 | import torchvision 11 | import torchvision.transforms as transforms 12 | 13 | import matplotlib.pyplot as plt 14 | 15 | import tensorboardX 16 | 17 | # Hyper parameters 18 | epoch_num = 200 19 | img_size = 64 # size of generated image 20 | batch_size = 128 21 | lr_g = 0.0001 # learning rate for Generator 22 | lr_c = 0.0001 # learning rate for Critic 23 | latent = 100 # dim of latent space 24 | img_channel = 1 # channel of generated image 25 | init_channel = 16 # control the initial Conv channel of the Generator and Discriminator 26 | workers = 2 # subprocess number for load the image 27 | k = 5 # train Critic K times and then train Generator one time 28 | dataset_size = 60000 # image number of your training set 29 | 30 | mean = [0.5] 31 | std = [0.5] 32 | 33 | slope = 0.2 # slope for leaky relu 34 | 35 | penalty_coeffcient = 10 # penalty coeffcient 36 | 37 | beta = (0.5, 0.9) # Adam hyperpaarmeters 38 | 39 | # use tensorboard 40 | writer = tensorboardX.SummaryWriter(log_dir='./logs/') 41 | 42 | # data enhancement 43 | data_transform = transforms.Compose([ 44 | transforms.Resize(size=img_size), 45 | transforms.ToTensor(), 46 | transforms.Normalize(mean=mean, std=std) 47 | ]) 48 | 49 | # dataset 50 | data_set = torchvision.datasets.MNIST( 51 | root='./MNIST', 52 | train=True, 53 | transform=data_transform 54 | ) 55 | 56 | data_loader = DataLoader(data_set, batch_size, True, num_workers=workers) 57 | 58 | 59 | # Generator 60 | class Generator(nn.Module): 61 | def __init__(self): 62 | super().__init__() 63 | 64 | self.deconv1 = nn.Sequential( 65 | nn.ConvTranspose2d(latent, init_channel * 8, 4, bias=False), 66 | nn.BatchNorm2d(init_channel * 8), 67 | nn.ReLU() 68 | ) 69 | 70 | self.deconv2 = nn.Sequential( 71 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 72 | nn.BatchNorm2d(init_channel * 4), 73 | nn.ReLU() 74 | ) 75 | 76 | self.deconv3 = nn.Sequential( 77 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 78 | nn.BatchNorm2d(init_channel * 2), 79 | nn.ReLU() 80 | ) 81 | 82 | self.deconv4 = nn.Sequential( 83 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 84 | nn.BatchNorm2d(init_channel), 85 | nn.ReLU() 86 | ) 87 | 88 | self.deconv5 = nn.Sequential( 89 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 90 | nn.Tanh() 91 | ) 92 | 93 | # initialization for parameters 94 | 95 | for layer in self.modules(): 96 | 97 | if isinstance(layer, nn.ConvTranspose2d): 98 | 99 | nn.init.normal(layer.weight.data, 0, 0.02) 100 | 101 | elif isinstance(layer, nn.BatchNorm2d): 102 | layer.weight.data.fill_(1) 103 | layer.bias.data.zero_() 104 | 105 | def forward(self, inputs): 106 | 107 | outputs = self.deconv1(inputs) 108 | outputs = self.deconv2(outputs) 109 | outputs = self.deconv3(outputs) 110 | outputs = self.deconv4(outputs) 111 | outputs = self.deconv5(outputs) 112 | 113 | return outputs 114 | 115 | 116 | # Critic : no BN!!! instead we can use layerNorm as suggest in the paper 117 | class Critic(nn.Module): 118 | def __init__(self): 119 | super().__init__() 120 | 121 | self.conv1 = nn.Sequential( 122 | nn.Conv2d(img_channel, init_channel, 4, 2, 1, bias=False), 123 | #nn.LayerNorm([init_channel, img_size // 2, img_size // 2]), 124 | nn.LeakyReLU(slope) 125 | ) 126 | 127 | self.conv2 = nn.Sequential( 128 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 129 | #nn.LayerNorm([init_channel * 2, img_size // 4, img_size // 4]), 130 | nn.LeakyReLU(slope) 131 | ) 132 | 133 | self.conv3 = nn.Sequential( 134 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 135 | #nn.LayerNorm([init_channel * 4, img_size // 8, img_size // 8]), 136 | nn.LeakyReLU(slope) 137 | ) 138 | 139 | self.conv4 = nn.Sequential( 140 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 141 | #nn.LayerNorm([init_channel * 8, img_size // 16, img_size // 16]), 142 | nn.LeakyReLU(slope) 143 | ) 144 | 145 | self.conv5 = nn.Sequential( 146 | 147 | # no sigmoid! 148 | nn.Conv2d(init_channel * 8, 1, 4, bias=False) 149 | ) 150 | 151 | # initialization for parameters 152 | for layer in self.modules(): 153 | 154 | if isinstance(layer, nn.Conv2d): 155 | nn.init.normal(layer.weight.data, 0, 0.02) 156 | 157 | def forward(self, inputs): 158 | 159 | outputs = self.conv1(inputs) 160 | outputs = self.conv2(outputs) 161 | outputs = self.conv3(outputs) 162 | outputs = self.conv4(outputs) 163 | outputs = self.conv5(outputs) 164 | 165 | return outputs.view(inputs.size(0)) 166 | 167 | 168 | # use cuda if you have GPU 169 | net_g = Generator().cuda() 170 | net_c = Critic().cuda() 171 | 172 | # optimizer 173 | opt_g = torch.optim.Adam(net_g.parameters(), lr=lr_g, betas=beta) # optimizer for Generator 174 | opt_c = torch.optim.Adam(net_c.parameters(), lr=lr_c, betas=beta) # optimizer for Critic 175 | 176 | 177 | # get random noise 178 | def get_noise(noise_num=batch_size): 179 | 180 | return Variable(torch.randn((noise_num, latent, 1, 1)).cuda()) 181 | 182 | 183 | # train the network 184 | start = time() 185 | fix_noise = get_noise(64) 186 | 187 | for epoch in range(epoch_num): 188 | 189 | for step, (real_data, target) in enumerate(data_loader, 1): 190 | 191 | # train Critic 192 | 193 | real_data = Variable(real_data).cuda() 194 | 195 | fake_data = net_g(get_noise(real_data.size(0))) 196 | 197 | output_fake = net_c(fake_data) 198 | 199 | output_real = net_c(real_data) 200 | 201 | # Calculate gradient penalty 202 | 203 | e = torch.rand(real_data.size(0), 1, 1, 1).cuda() 204 | 205 | interpolation = Variable(e * real_data.data + (1 - e) * fake_data.data, requires_grad=True) 206 | 207 | gradient_penalty = torch.autograd.grad(outputs=net_c(interpolation), inputs=interpolation, 208 | grad_outputs=torch.ones(net_c(interpolation).size()).cuda(), 209 | create_graph=True, retain_graph=True)[0] 210 | 211 | opt_c.zero_grad() 212 | 213 | gradient_penalty = gradient_penalty.view(gradient_penalty.size(0), -1) 214 | 215 | loss_c = torch.mean(output_fake - output_real + penalty_coeffcient * (gradient_penalty.norm(2, dim=1) - 1) ** 2) 216 | 217 | loss_c.backward() 218 | opt_c.step() 219 | 220 | # train Generator 221 | if step % k is 0: 222 | 223 | output_fake = net_c(net_g(get_noise())) 224 | 225 | loss_g = - torch.mean(output_fake) 226 | 227 | opt_g.zero_grad() 228 | loss_g.backward() 229 | opt_g.step() 230 | 231 | if step % 20 is 0: 232 | 233 | iteration = epoch * (math.ceil(dataset_size / batch_size)) + step 234 | writer.add_scalars('loss', {'g_loss': loss_g.data.item(), 'c_loss': loss_c.data.item()}, int(iteration)) 235 | print('epoch:', epoch, 'step', step, 'time:', (time() - start) / 60, 'min') 236 | 237 | fake_img = torchvision.utils.make_grid((0.5 * net_g(fix_noise).data.cpu() + 0.5)) 238 | plt.imshow(fake_img.permute(1, 2, 0).numpy()) 239 | plt.pause(0.01) 240 | 241 | plt.show() 242 | -------------------------------------------------------------------------------- /GANs/WGAN-GP-humanface.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import torch 4 | import torch.nn as nn 5 | from torch.utils.data import DataLoader 6 | from torch.autograd import Variable 7 | 8 | import torchvision 9 | import torchvision.transforms as transforms 10 | 11 | import cv2 as cv 12 | 13 | torch.cuda.set_device(1) 14 | 15 | # Hyper parameters 16 | epoch_num = 400 17 | img_size = 64 # size of generated image 18 | batch_size = 128 19 | lr_g = 0.0001 # learning rate for Generator 20 | lr_c = 0.0001 # learning rate for Critic 21 | latent = 100 # dim of latent space 22 | img_channel = 3 # channel of generated image 23 | init_channel = 64 # control the initial Conv channel of the Generator and Discriminator 24 | workers = 2 # subprocess number for load the image 25 | k = 5 # train Critic K times and then train Generator one time 26 | 27 | mean = [0.5] 28 | std = [0.5] 29 | 30 | slope = 0.2 # slope for leaky relu 31 | 32 | penalty_coeffcient = 10 # penalty coeffcient 33 | 34 | beta = (0.5, 0.9) # Adam hyperpaarmeters 35 | 36 | # data enhancement 37 | data_transform = transforms.Compose([ 38 | transforms.Resize(size=img_size), 39 | transforms.RandomHorizontalFlip(), 40 | transforms.ToTensor(), 41 | transforms.Normalize(mean=mean, std=std) 42 | ]) 43 | 44 | # dataset 45 | data_set = torchvision.datasets.ImageFolder(root='./data/', transform=data_transform) 46 | 47 | data_loader = DataLoader(data_set, batch_size, True, num_workers=workers) 48 | 49 | 50 | # Generator 51 | class Generator(nn.Module): 52 | def __init__(self): 53 | super().__init__() 54 | 55 | self.deconv1 = nn.Sequential( 56 | nn.ConvTranspose2d(latent, init_channel * 8, 4, bias=False), 57 | nn.BatchNorm2d(init_channel * 8), 58 | nn.ReLU(), 59 | 60 | nn.Conv2d(init_channel * 8, init_channel * 8, 3, padding=1, bias=False), 61 | nn.BatchNorm2d(init_channel * 8), 62 | nn.ReLU() 63 | ) 64 | 65 | self.deconv2 = nn.Sequential( 66 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 67 | nn.BatchNorm2d(init_channel * 4), 68 | nn.ReLU(), 69 | 70 | nn.Conv2d(init_channel * 4, init_channel * 4, 3, padding=1, bias=False), 71 | nn.BatchNorm2d(init_channel * 4), 72 | nn.ReLU() 73 | ) 74 | 75 | self.deconv3 = nn.Sequential( 76 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 77 | nn.BatchNorm2d(init_channel * 2), 78 | nn.ReLU(), 79 | 80 | nn.Conv2d(init_channel * 2, init_channel * 2, 3, padding=1, bias=False), 81 | nn.BatchNorm2d(init_channel * 2), 82 | nn.ReLU() 83 | ) 84 | 85 | self.deconv4 = nn.Sequential( 86 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 87 | nn.BatchNorm2d(init_channel), 88 | nn.ReLU(), 89 | 90 | nn.Conv2d(init_channel, init_channel, 3, padding=1, bias=False), 91 | nn.BatchNorm2d(init_channel), 92 | nn.ReLU() 93 | ) 94 | 95 | self.deconv5 = nn.Sequential( 96 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 97 | nn.Tanh() 98 | ) 99 | 100 | # initialization for parameters 101 | 102 | for layer in self.modules(): 103 | 104 | if isinstance(layer, nn.ConvTranspose2d): 105 | 106 | nn.init.normal(layer.weight.data, 0, 0.02) 107 | 108 | elif isinstance(layer, nn.BatchNorm2d): 109 | layer.weight.data.fill_(1) 110 | layer.bias.data.zero_() 111 | 112 | def forward(self, inputs): 113 | 114 | outputs = self.deconv1(inputs) 115 | outputs = self.deconv2(outputs) 116 | outputs = self.deconv3(outputs) 117 | outputs = self.deconv4(outputs) 118 | outputs = self.deconv5(outputs) 119 | 120 | return outputs 121 | 122 | 123 | # Critic : no BN!!! instead we can use layerNorm as suggest in the paper 124 | class Critic(nn.Module): 125 | def __init__(self): 126 | super().__init__() 127 | 128 | self.conv1 = nn.Sequential( 129 | nn.Conv2d(img_channel, init_channel, 4, 2, 1, bias=False), 130 | nn.LeakyReLU(slope), 131 | 132 | nn.Conv2d(init_channel, init_channel, 3, padding=1), 133 | nn.LeakyReLU(slope) 134 | ) 135 | 136 | self.conv2 = nn.Sequential( 137 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 138 | nn.LeakyReLU(slope), 139 | 140 | nn.Conv2d(init_channel * 2, init_channel * 2, 3, padding=1), 141 | nn.LeakyReLU(slope) 142 | 143 | ) 144 | 145 | self.conv3 = nn.Sequential( 146 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 147 | nn.LeakyReLU(slope), 148 | 149 | nn.Conv2d(init_channel * 4, init_channel * 4, 3, padding=1), 150 | nn.LeakyReLU(slope) 151 | 152 | ) 153 | 154 | self.conv4 = nn.Sequential( 155 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 156 | nn.LeakyReLU(slope), 157 | 158 | nn.Conv2d(init_channel * 8, init_channel * 8, 3, padding=1), 159 | nn.LeakyReLU(slope) 160 | 161 | ) 162 | 163 | self.conv5 = nn.Sequential( 164 | 165 | # no sigmoid! 166 | nn.Conv2d(init_channel * 8, 1, 4, bias=False) 167 | ) 168 | 169 | # initialization for parameters 170 | for layer in self.modules(): 171 | 172 | if isinstance(layer, nn.Conv2d): 173 | nn.init.normal(layer.weight.data, 0, 0.02) 174 | 175 | def forward(self, inputs): 176 | 177 | outputs = self.conv1(inputs) 178 | outputs = self.conv2(outputs) 179 | outputs = self.conv3(outputs) 180 | outputs = self.conv4(outputs) 181 | outputs = self.conv5(outputs) 182 | 183 | return outputs.view(inputs.size(0)) 184 | 185 | 186 | # use cuda if you have GPU 187 | net_g = Generator().cuda() 188 | net_c = Critic().cuda() 189 | 190 | # optimizer 191 | opt_g = torch.optim.Adam(net_g.parameters(), lr=lr_g, betas=beta) # optimizer for Generator 192 | opt_c = torch.optim.Adam(net_c.parameters(), lr=lr_c, betas=beta) # optimizer for Critic 193 | 194 | 195 | # get random noise 196 | def get_noise(noise_num=batch_size): 197 | 198 | return Variable(torch.randn((noise_num, latent, 1, 1)).cuda()) 199 | 200 | 201 | # train the network 202 | start = time() 203 | number = 1 204 | 205 | for epoch in range(epoch_num): 206 | 207 | for step, (real_data, target) in enumerate(data_loader, 1): 208 | 209 | # train Critic 210 | 211 | real_data = Variable(real_data).cuda() 212 | 213 | fake_data = net_g(get_noise(real_data.size(0))) 214 | 215 | output_fake = net_c(fake_data) 216 | 217 | output_real = net_c(real_data) 218 | 219 | # Calculate gradient penalty 220 | 221 | e = torch.rand(real_data.size(0), 1, 1, 1).cuda() 222 | 223 | interpolation = Variable(e * real_data.data + (1 - e) * fake_data.data, requires_grad=True) 224 | 225 | gradient_penalty = torch.autograd.grad(outputs=net_c(interpolation), inputs=interpolation, 226 | grad_outputs=torch.ones(net_c(interpolation).size()).cuda(), 227 | create_graph=True, retain_graph=True)[0] 228 | 229 | opt_c.zero_grad() 230 | 231 | gradient_penalty = gradient_penalty.view(gradient_penalty.size(0), -1) 232 | 233 | loss_c = torch.mean(output_fake - output_real + penalty_coeffcient * (gradient_penalty.norm(2, dim=1) - 1) ** 2) 234 | 235 | loss_c.backward() 236 | opt_c.step() 237 | 238 | # train Generator 239 | if step % k is 0: 240 | 241 | output_fake = net_c(net_g(get_noise())) 242 | 243 | loss_g = - torch.mean(output_fake) 244 | 245 | opt_g.zero_grad() 246 | loss_g.backward() 247 | opt_g.step() 248 | 249 | if step % 20 is 0: 250 | 251 | fake_img = torchvision.utils.make_grid((0.5 * net_g(get_noise(64)).data.cpu() + 0.5)) 252 | cv.imwrite('./img/WGAN/' + str(number) + '.png', fake_img.permute(1, 2, 0).numpy() * 255) 253 | 254 | number += 1 255 | 256 | 257 | 258 | -------------------------------------------------------------------------------- /GANs/CGAN-MNIST.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import torch 4 | import torch.nn as nn 5 | from torch.utils.data import DataLoader 6 | from torch.autograd import Variable 7 | 8 | import torchvision 9 | import torchvision.transforms as transforms 10 | 11 | import matplotlib.pyplot as plt 12 | 13 | 14 | # Hyper parameters 15 | epoch_num = 200 16 | img_size = 64 # size of generated image 17 | batch_size = 128 18 | lr_g = 0.0002 # learning rate for Generator 19 | lr_d = 0.0002 # learning rate for Discriminator 20 | latent = 100 # dim of latent space 21 | conditional_code = 10 # dim of conditional code 22 | img_channel = 1 # channel of generated image 23 | init_channel = 16 # control the initial Conv channel of the Generator and Discriminator 24 | workers = 1 # subprocess number for load the image 25 | k = 2 # train Discriminator K times and then train Generator one time 26 | 27 | mean = [0.5] 28 | std = [0.5] 29 | 30 | slope = 0.2 # slope for leaky relu 31 | 32 | # data enhancement 33 | data_transform = transforms.Compose([ 34 | transforms.Resize(size=img_size), 35 | transforms.ToTensor(), 36 | transforms.Normalize(mean=mean, std=std) 37 | ]) 38 | 39 | # dataset 40 | data_set = torchvision.datasets.MNIST('./MNIST/', transform=data_transform) 41 | 42 | data_loader = DataLoader(data_set, batch_size, True, num_workers=workers) 43 | 44 | 45 | # Generator 46 | class Generator(nn.Module): 47 | def __init__(self): 48 | super().__init__() 49 | 50 | self.deconv1_1 = nn.Sequential( # input for latent coda 51 | nn.ConvTranspose2d(latent, init_channel * 4, 4, bias=False), 52 | nn.BatchNorm2d(init_channel * 4), 53 | nn.ReLU(), 54 | ) 55 | 56 | self.deconv1_2 = nn.Sequential( # input for conditional code 57 | nn.ConvTranspose2d(conditional_code, init_channel * 4, 4, bias=False), 58 | nn.BatchNorm2d(init_channel * 4), 59 | nn.ReLU(), 60 | ) 61 | 62 | self.deconv2 = nn.Sequential( 63 | nn.ConvTranspose2d(init_channel * 8, init_channel * 4, 4, 2, 1, bias=False), 64 | nn.BatchNorm2d(init_channel * 4), 65 | nn.ReLU(), 66 | ) 67 | 68 | self.deconv3 = nn.Sequential( 69 | nn.ConvTranspose2d(init_channel * 4, init_channel * 2, 4, 2, 1, bias=False), 70 | nn.BatchNorm2d(init_channel * 2), 71 | nn.ReLU(), 72 | ) 73 | 74 | self.deconv4 = nn.Sequential( 75 | nn.ConvTranspose2d(init_channel * 2, init_channel, 4, 2, 1, bias=False), 76 | nn.BatchNorm2d(init_channel), 77 | nn.ReLU(), 78 | ) 79 | 80 | self.deconv5 = nn.Sequential( 81 | nn.ConvTranspose2d(init_channel, img_channel, 4, 2, 1, bias=False), 82 | nn.Tanh() 83 | ) 84 | 85 | # initialization for parameters 86 | 87 | for layer in self.modules(): 88 | 89 | if isinstance(layer, nn.ConvTranspose2d): 90 | nn.init.normal(layer.weight.data, 0, 0.02) 91 | 92 | elif isinstance(layer, nn.BatchNorm2d): 93 | layer.weight.data.fill_(1) 94 | layer.bias.data.zero_() 95 | 96 | def forward(self, inputs, labels): 97 | 98 | outputs = torch.cat((self.deconv1_1(inputs), self.deconv1_2(labels)), dim=1) 99 | 100 | outputs = self.deconv2(outputs) 101 | outputs = self.deconv3(outputs) 102 | outputs = self.deconv4(outputs) 103 | outputs = self.deconv5(outputs) 104 | 105 | return outputs 106 | 107 | 108 | # Discriminator( 109 | class Discriminator(nn.Module): 110 | def __init__(self): 111 | super().__init__() 112 | 113 | self.conv1_1 = nn.Sequential( # input for image 114 | nn.Conv2d(img_channel, init_channel // 2, 4, 2, 1, bias=False), 115 | nn.LeakyReLU(slope) 116 | ) 117 | 118 | self.conv1_2 = nn.Sequential( # input for onehot code 119 | nn.Conv2d(conditional_code, init_channel // 2, 4, 2, 1, bias=False), 120 | nn.LeakyReLU(slope) 121 | ) 122 | 123 | self.conv2 = nn.Sequential( 124 | nn.Conv2d(init_channel, init_channel * 2, 4, 2, 1, bias=False), 125 | nn.BatchNorm2d(init_channel * 2), 126 | nn.LeakyReLU(slope) 127 | ) 128 | 129 | self.conv3 = nn.Sequential( 130 | nn.Conv2d(init_channel * 2, init_channel * 4, 4, 2, 1, bias=False), 131 | nn.BatchNorm2d(init_channel * 4), 132 | nn.LeakyReLU(slope) 133 | ) 134 | 135 | self.conv4 = nn.Sequential( 136 | nn.Conv2d(init_channel * 4, init_channel * 8, 4, 2, 1, bias=False), 137 | nn.BatchNorm2d(init_channel * 8), 138 | nn.LeakyReLU(slope) 139 | ) 140 | 141 | self.conv5 = nn.Sequential( 142 | nn.Conv2d(init_channel * 8, 1, 4, bias=False), 143 | nn.Sigmoid() 144 | ) 145 | 146 | # initialization for parameters 147 | for layer in self.modules(): 148 | 149 | if isinstance(layer, nn.Conv2d): 150 | nn.init.normal(layer.weight.data, 0, 0.02) 151 | 152 | elif isinstance(layer, nn.BatchNorm2d): 153 | layer.weight.data.fill_(1) 154 | layer.bias.data.zero_() 155 | 156 | def forward(self, inputs, labels): 157 | 158 | outputs = torch.cat((self.conv1_1(inputs), self.conv1_2(labels)), dim=1) 159 | 160 | outputs = self.conv2(outputs) 161 | outputs = self.conv3(outputs) 162 | outputs = self.conv4(outputs) 163 | outputs = self.conv5(outputs) 164 | 165 | return outputs.view(inputs.size(0)) 166 | 167 | 168 | # use cuda if you have GPU 169 | net_g = Generator().cuda() 170 | net_d = Discriminator().cuda() 171 | 172 | # optimizer 173 | opt_g = torch.optim.Adam(net_g.parameters(), lr=lr_g, betas=(0.5, 0.999)) # optimizer for Generator 174 | opt_d = torch.optim.Adam(net_d.parameters(), lr=lr_d, betas=(0.5, 0.999)) # optimizer for Discriminator 175 | 176 | 177 | # get random noise 178 | def get_noise(noise_num=batch_size): 179 | 180 | return Variable(torch.randn((noise_num, latent, 1, 1)).cuda()) 181 | 182 | 183 | # transform a number to one hot code 184 | def one_hot(number): 185 | 186 | temp = torch.zeros((number.size(0), conditional_code)) 187 | 188 | for i in range(number.size(0)): 189 | 190 | temp[i][int(number[i])] = 1 191 | 192 | temp = temp.view(temp.size(0), temp.size(1), 1, 1) 193 | 194 | return Variable(temp).cuda() 195 | 196 | 197 | # expand a onehot code to image size 198 | def one_hot_expand(onehot_code): 199 | 200 | return onehot_code.expand(onehot_code.size(0), onehot_code.size(1), img_size, img_size) 201 | 202 | 203 | # before train the network 204 | start = time() 205 | number = 1 206 | 207 | labellist = [] 208 | 209 | for i in range(10): 210 | 211 | for j in range(8): 212 | 213 | labellist.append(i) 214 | 215 | 216 | # train the network 217 | for epoch in range(epoch_num): 218 | 219 | for step, (real_data, target) in enumerate(data_loader, 1): 220 | 221 | # train Discriminator 222 | 223 | real_data = Variable(real_data).cuda() 224 | 225 | real_label = one_hot(target) 226 | 227 | fake_label = one_hot(torch.floor(torch.rand(real_data.size(0)) * conditional_code)) 228 | 229 | prob_fake = net_d(net_g(get_noise(real_data.size(0)), fake_label), one_hot_expand(fake_label)) 230 | 231 | prob_real = net_d(real_data, one_hot_expand(real_label)) 232 | 233 | loss_d = - torch.mean(torch.log(prob_real) + torch.log(1 - prob_fake)) 234 | 235 | opt_d.zero_grad() 236 | loss_d.backward() 237 | opt_d.step() 238 | 239 | # train Generator 240 | if step % k is 0: 241 | 242 | fake_label = one_hot(torch.floor(torch.rand(batch_size) * conditional_code)) 243 | 244 | prob_fake = net_d(net_g(get_noise(), fake_label), one_hot_expand(fake_label)) 245 | 246 | loss_g = - torch.mean(torch.log(prob_fake)) 247 | 248 | opt_g.zero_grad() 249 | loss_g.backward() 250 | opt_g.step() 251 | 252 | if step % 20 is 0: 253 | 254 | print('epoch:', epoch, 'step', step, 'time:', (time() - start) / 60, 'min') 255 | 256 | fake_label = one_hot(torch.FloatTensor(labellist)) 257 | 258 | generate_img = torchvision.utils.make_grid((0.5 * net_g(get_noise(80), fake_label).data.cpu() + 0.5)) 259 | 260 | plt.imshow(generate_img.permute(1, 2, 0).numpy()) 261 | plt.pause(0.01) 262 | 263 | plt.imsave('./img/' + str(number) + '.png', generate_img.permute(1, 2, 0).numpy()) 264 | 265 | number += 1 266 | 267 | --------------------------------------------------------------------------------