Conditional Gan代码学习_vaegan中的normal_init函数

 自己修改了github的代码，选取了比较核心的代码，简化了相关操作

 基于MNIST数据集

# 库的引用
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy
from torchvision import datasets
from torch.utils.data import DataLoader
from torchvision import transforms
from torch.autograd import Variable
from torchvision.utils import save_image
 
 
# 用到的参数及设置
batch_size = 128
lr = 0.0002
train_epoch = 50
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
 
 
# 准备数据
data = datasets.MNIST(root='./data/mnist', train=True, transform=transform, download=True)
data_loader = DataLoader(data, batch_size=batch_size, shuffle=True, num_workers=4)
 
def normal_init(m, mean, std):
    if isinstance(m, nn.Linear):
        m.weight.data.normal_(mean, std)
        m.bias.data.zero_()
 
 
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.fc1_1 = nn.Linear(100, 256) #　100是噪声的维度
        self.fc1_1_bn = nn.BatchNorm1d(256)
        self.fc1_2 = nn.Linear(10, 256)
        self.fc1_2_bn = nn.BatchNorm1d(256)
        self.fc2 = nn.Linear(512, 512)
        self.fc2_bn = nn.BatchNorm1d(512)
        self.fc3 = nn.Linear(512, 1024)
        self.fc3_bn = nn.BatchNorm1d(1024)
        self.fc4 = nn.Linear(1024, 784)
 
 
    # 循环时字典遍历的是key
    def weight_init(self, mean, std):
        for m in self._modules:
            normal_init(self._modules[m], mean, std)
 
 
    def forward(self, input, label):
        batch_size = input.size(0)
        x = F.relu(self.fc1_1_bn(self.fc1_1(input)))
        y = F.relu(self.fc1_2_bn(self.fc1_2(label)))
        x = torch.cat([x, y], 1) # 这样就变为512维了
        x = F.relu(self.fc2_bn(self.fc2(x)))
        x = F.relu(self.fc3_bn(self.fc3(x)))
        x = torch.tanh(self.fc4(x))
        # 可以reshape为图片的维度
        gen_img = x.view(batch_size, 1, 28, 28)
        return gen_img
 
 
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.fc1_1 = nn.Linear(784, 1024)
        self.fc1_2 = nn.Linear(10, 1024)
        self.fc2 = nn.Linear(2048, 512)
        self.fc2_bn = nn.BatchNorm1d(512)
        self.fc3 = nn.Linear(512, 256)
        self.fc3_bn = nn.BatchNorm1d(256)
        self.fc4 = nn.Linear(256, 1)
 
    def weight_init(self, mean, std):
        for m in self._modules:
            normal_init(self._modules[m], mean, std)
 
    def forward(self, input, label):
        # Flatten操作，展平为784个像素后再操作
        batch_size = input.shape[0]
        x = input.view(batch_size, -1)
 
        x = F.leaky_relu(self.fc1_1(x), 0.2)
        y = F.leaky_relu(self.fc1_2(label), 0.2)
        x = torch.cat([x, y], 1) # 按列进行连接，变为2048维
        x = F.leaky_relu(self.fc2_bn(self.fc2(x)), 0.2)
        x = F.leaky_relu(self.fc3_bn(self.fc3(x)), 0.2)
        x = torch.sigmoid(self.fc4(x))
 
        return x
 
 
# network，结构可以参照图片
G = Generator().cuda()
D = Discriminator().cuda()
G.weight_init(mean=0, std=0.02)
D.weight_init(mean=0, std=0.02)
criterion = nn.BCELoss()
G_optimizer = torch.optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))
D_optimizer = torch.optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))
 
 
def train(epoch):
 
    # 随着训练进行， 让lr变小
 
    if (epoch+1) == 30:
        G_optimizer.param_groups[0]['lr'] /= 10
        D_optimizer.param_groups[0]['lr'] /= 10
        print("learning rate change!")
 
 
    if (epoch+1) == 40:
        G_optimizer.param_groups[0]['lr'] /= 10
        D_optimizer.param_groups[0]['lr'] /= 10
        print("learning rate change!")
 
 
    for ind, (images, labels) in enumerate(data_loader):
 
        # 随机噪声z，100维
        z = Variable(torch.Tensor(numpy.random.normal(0, 1, (images.size(0), 100))))
 
        # train discriminator D
 
        mini_batch = images.size(0)
        D.zero_grad()
 
        #这是对应的二分类真与假的标签值
        y_real_ = torch.ones(mini_batch)
        y_fake_ = torch.zeros(mini_batch)
 
        # 这两句是生成label的one-hot向量,因为bceloss用的是one-hot
        y_label_ = torch.zeros(mini_batch, 10)
        y_label_.scatter_(1, labels.view(mini_batch, 1), 1)
        # (128,10)
 
        images, y_label_, y_real_, y_fake_ = Variable(images.cuda()), Variable(y_label_.cuda()), Variable(\
            y_real_.cuda()), Variable(y_fake_.cuda())
 
 
        D_result = D(images, y_label_).squeeze()  # 这是告诉判别器标签和图片的配对情况
        # 这个地方suqeeze就是128， 不是就是(128,1)， 不能BCEloss
 
        D_real_loss = criterion(D_result, y_real_)  # Discriminator的一个loss
 
        y_ = (torch.rand(mini_batch, 1) * 10).type(torch.LongTensor)
        y_label_ = torch.zeros(mini_batch, 10)
        y_label_.scatter_(1, y_.view(mini_batch, 1), 1)
        z, y_label_ = Variable(z.cuda()), Variable(y_label_.cuda())
 
        G_result = G(z, y_label_)
        gen_imgs = G(z, y_label_)
 
        D_result = D(G_result, y_label_).squeeze()
        D_fake_loss = criterion(D_result, y_fake_)  # Discriminator的假图片loss
 
        D_train_loss = D_real_loss + D_fake_loss
 
        D_train_loss.backward()
        D_optimizer.step()
 
        # train generator G
        G.zero_grad()
 
        # 杂乱的生成的标签
        z_ = torch.rand((mini_batch, 100))
        y_ = (torch.rand(mini_batch, 1) * 10).type(torch.LongTensor)
        y_label_ = torch.zeros(mini_batch, 10)
        y_label_.scatter_(1, y_.view(mini_batch, 1), 1)
 
        z_, y_label_ = Variable(z_.cuda()), Variable(y_label_.cuda())
 
        G_result = G(z_, y_label_)
        D_result = D(G_result, y_label_).squeeze()
        G_train_loss = criterion(D_result, y_real_)
        G_train_loss.backward()
        G_optimizer.step()
 
 
        if ind % 50 == 0:
            print('G_loss:', G_train_loss.data.item(), 'D_loss:', D_train_loss.data.item())
    save_image(gen_imgs.data[:25], 'images{}.png'.format(epoch), nrow=5, normalize=True)
 
 
if __name__ == '__main__':
    for epoch in range(train_epoch):
        train(epoch)

 结果如图，效果还是不错的