【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)
相关推荐
zzywxc7872 分钟前
AI大模型的技术演进、流程重构、行业影响三个维度的系统性分析
人工智能·重构
点控云2 分钟前
智能私域运营中枢:从客户视角看 SCRM 的体验革新与价值重构
大数据·人工智能·科技·重构·外呼系统·呼叫中心
zhaoyi_he10 分钟前
多模态大模型的技术应用与未来展望:重构AI交互范式的新引擎
人工智能·重构
葫三生1 小时前
如何评价《论三生原理》在科技界的地位?
人工智能·算法·机器学习·数学建模·量子计算
m0_751336392 小时前
突破性进展:超短等离子体脉冲实现单电子量子干涉,为飞行量子比特奠定基础
人工智能·深度学习·量子计算·材料科学·光子器件·光子学·无线电电子
美狐美颜sdk5 小时前
跨平台直播美颜SDK集成实录:Android/iOS如何适配贴纸功能
android·人工智能·ios·架构·音视频·美颜sdk·第三方美颜sdk
DeepSeek-大模型系统教程6 小时前
推荐 7 个本周 yyds 的 GitHub 项目。
人工智能·ai·语言模型·大模型·github·ai大模型·大模型学习
郭庆汝6 小时前
pytorch、torchvision与python版本对应关系
人工智能·pytorch·python
小雷FansUnion8 小时前
深入理解MCP架构:智能服务编排、上下文管理与动态路由实战
人工智能·架构·大模型·mcp
资讯分享周8 小时前
扣子空间PPT生产力升级:AI智能生成与多模态创作新时代
人工智能·powerpoint