【GAN】简单的GAN模型搭建 -- 以线性模型和MNIST数据集为例子

文章目录

不讲原理,从简单的代码一步步开始,学会怎么用、怎么设计损失函数即可。

确定损失函数

生成器的任务是生成足够以假乱真的数据,判别器的任务是分辨出哪些数据是真实的,哪些数据是假的。因此,对于判别器来讲,需要判别真伪,也就是true/false,从这个角度看,是个二分类问题。所以损失函数使用二类分类损失,即BCELoss。

python 复制代码
import torch
import torch.nn as nn

adversarial_loss = nn.BCELoss()

生成器网络架构

这里使用纯线性网络作为生成器,得到的输出为[batch_size, np.pord(28*28)]

python 复制代码
import torch.nn as nn
import numpy as np

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

        def block(in_features, out_features, normalization=True):
            layers = [nn.Linear(in_features, out_features)]
            if normalization:
                layers.append(nn.BatchNorm1d(out_features, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers

        self.model = nn.Sequential(
            *block(100, 128, normalization=False),
            *block(128, 256),
            *block(256, 512),
            *block(512, 1024),
            nn.Linear(1024, int(np.prod((1, 28, 28)))), # generate a photo size, but in line mode.
            nn.Tanh()
        )

    def forward(self, x):
        return self.model(x)

判别器网络架构:

判别器的功能为判断出来哪一个是生成的图片,哪一个是真实的图片。对于生成的图片,我们希望判别器打上假的标签,对于真实的图片,我们希望判别器打上真的标签,因此,判别器的输出为一个数,即0或者1。

python 复制代码
import torch.nn as nn
import numpy as np
class Disctiminator(nn.Module):
    def __init__(self):
        super(Disctiminator, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(int(np.prod((1, 28, 28))), 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        img_flat = x.view(x.size(0), -1)
        validity = self.model(img_flat)
        return validity

训练流程:

载入数据 --- 训练 (生成图片 --- 损失 --- 反向传播) --- 测试(这里没有加测试代码,可以照着训练代码改一下)

损失函数:生成器损失函数和判别器损失函数,两个损失函数分别进行反向传播,即生成器损失函数优化生成器,判别器损失函数优化判别器。

python 复制代码
import torch
import torch.nn as nn
import argparse
import os
import numpy as np
from torch.utils.data import DataLoader, Dataset, dataset
from torchvision import datasets
from torchvision.transforms import transforms
from torchvision.utils import save_image
from torch.autograd import Variable
from models.generator import Generator
from models.dicsriminator import Disctiminator





os.makedirs('/home/sjr/gxj/study/data/mnist', exist_ok=True)

dataloader = DataLoader(datasets.MNIST('/home/sjr/gxj/study/data/mnist',
                                       train=True, download=True,
                                       transform=transforms.Compose([transforms.Resize(28), transforms.ToTensor(),
                                                                     transforms.Normalize([0.5], [0.5])])
                                       ),
                        batch_size=64, shuffle=True, num_workers=4)

adversarial_loss = torch.nn.BCELoss()
generator = Generator()
discriminator = Disctiminator()

device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
if torch.cuda.is_available():
    adversarial_loss.cuda(device)
    generator.cuda(device)
    discriminator.cuda(device)

optimizer_G = torch.optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))

Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor

for epoch in range(60):
    for i, (imgs, _) in enumerate(dataloader):
        valid = Tensor(imgs.size(0), 1).fill_(1.0)
        fake = Tensor(imgs.size(0), 1).fill_(0.0)
        real_imgs = Variable(imgs.type(Tensor))

        optimizer_G.zero_grad()

        z = Tensor(np.random.normal(0, 1, (imgs.shape[0], 100)))
        gen_imgs = generator(z)


        g_loss = adversarial_loss(discriminator(gen_imgs), valid)
        g_loss.backward()
        optimizer_G.step()

        optimizer_D.zero_grad()

        real_loss = adversarial_loss(discriminator(real_imgs), valid)
        fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)
        d_loss = (real_loss + fake_loss) / 2
        d_loss.backward()
        optimizer_D.step()
        #
        print(f"[Epoch {epoch}/{200}] [Batch {i}/{len(dataloader)}] [D loss: {d_loss.item()}] [G loss: {g_loss.item()}]")
        #
        if (epoch + 1) % 20 == 0:
            save_image(gen_imgs.data[:25], f'images/{epoch+1}.png', nrow=5, normalize=True)
相关推荐
机器人虎哥几秒前
【8210A-TX2】Ubuntu18.04 + ROS_ Melodic + TM-16多线激光 雷达评测
人工智能·机器学习
码银8 分钟前
冲破AI 浪潮冲击下的 迷茫与焦虑
人工智能
何大春11 分钟前
【弱监督语义分割】Self-supervised Image-specific Prototype Exploration for WSSS 论文阅读
论文阅读·人工智能·python·深度学习·论文笔记·原型模式
uncle_ll19 分钟前
PyTorch图像预处理:计算均值和方差以实现标准化
图像处理·人工智能·pytorch·均值算法·标准化
宋1381027972020 分钟前
Manus Xsens Metagloves虚拟现实手套
人工智能·机器人·vr·动作捕捉
在下不上天20 分钟前
Flume日志采集系统的部署,实现flume负载均衡,flume故障恢复
大数据·开发语言·python
SEVEN-YEARS23 分钟前
深入理解TensorFlow中的形状处理函数
人工智能·python·tensorflow
世优科技虚拟人27 分钟前
AI、VR与空间计算:教育和文旅领域的数字转型力量
人工智能·vr·空间计算
EterNity_TiMe_28 分钟前
【论文复现】(CLIP)文本也能和图像配对
python·学习·算法·性能优化·数据分析·clip
cloud studio AI应用33 分钟前
腾讯云 AI 代码助手:产品研发过程的思考和方法论
人工智能·云计算·腾讯云