CGAN模型——pytorch实现

论文传送门：Conditional Generative Adversarial Nets 

import torch
import torch.nn as nn


class Discriminator(nn.Module):  # 定义判别器
    def __init__(self, img_size=(28, 28), num_classes=2):  # 初始化方法
        super(Discriminator, self).__init__()  # 继承初始化方法
        self.img_size = img_size  # 图片尺寸，默认单通道灰度图
        self.num_classes = num_classes  # 类别数

        self.linear1 = nn.Linear(self.img_size[0] * self.img_size[1] + self.num_classes, 512)  # linear映射
        self.linear2 = nn.Linear(512, 512)  # linear映射
        self.linear3 = nn.Linear(512, 512)  # linear映射
        self.linear4 = nn.Linear(512, 1)  # linear映射
        self.linear5 = nn.Linear(512, self.num_classes)  # linear映射
        self.dropout = nn.Dropout(0.4)  # dropout操作
        self.leakyrelu = nn.LeakyReLU(0.2, inplace=True)  # leakyrelu激活函数
        self.sigmoid = nn.Sigmoid()  # sigmoid激活函数
        self.softmax = nn.Softmax(dim=1)  # softmax激活函数
        self.embedding = nn.Embedding(self.num_classes, self.num_classes)  # embedding操作

    def forward(self, input: tuple):  # 前传函数
        image, label = input  # 从输入的元组中获取噪声向量和标签信息
        label = self.embedding(label)  # 标签信息经过embedding操作
        x = torch.flatten(image, 1)  # 输入图片从三维压缩至一维特征向量，(n,1,28,28)-->(n,784)
        x = torch.cat([x, label], dim=-1)  # 特征向量与标签向量进行拼接，(n,784)cat(n,2)-->(n,786)
        x = self.linear1(x)  # linear映射，(n,786)-->(n,512)
        x = self.leakyrelu(x)  # leakyrelu激活函数
        x = self.linear2(x)  # linear映射,(n,512)-->(n,512)
        x = self.leakyrelu(x)  # leakyrelu激活函数
        x = self.dropout(x)  # dropout操作
        x = self.linear3(x)  # linear映射,(n,512)-->(n,512)
        x = self.leakyrelu(x)  # leakyrelu激活函数
        x = self.dropout(x)  # dropout操作
        # 根据特征向量x，计算图片真假的得分
        validity = self.linear4(x)  # linear映射,(n,512)-->(n,1)
        validity = self.sigmoid(validity)  # sigmoid激活函数，将输出压缩至（0，1）
        # 根据特征向量x，计算图片分类的标签
        label = self.linear5(x)  # linear映射,(n,512)-->(n,2)
        label = self.softmax(label)  # softmax激活函数，将输出压缩至（0，1）

        return (validity, label)  # 返回(图片真假的得分，图片分类的标签)


class Generator(nn.Module):  # 定义生成器
    def __init__(self, img_size=(28, 28), num_classes=2, latent_dim=100):  # 初始化方法
        super(Generator, self).__init__()
        self.img_size = img_size  # 图片尺寸，默认单通道灰度图
        self.num_classes = num_classes  # 类别数
        self.latent_dim = latent_dim  # 噪声z的长度

        self.linear1 = nn.Linear(self.latent_dim + self.num_classes, 256)  # linear映射
        self.bn1 = nn.BatchNorm1d(256, 0.8)  # bn操作
        self.linear2 = nn.Linear(256, 512)  # linear映射
        self.bn2 = nn.BatchNorm1d(512, 0.8)  # bn操作
        self.linear3 = nn.Linear(512, 1024)  # linear映射
        self.bn3 = nn.BatchNorm1d(1024, 0.8)  # bn操作
        self.linear4 = nn.Linear(1024, self.img_size[0] * self.img_size[1])  # linear映射
        self.leakyrelu = nn.LeakyReLU(0.2, inplace=True)  # leakyrelu激活函数
        self.tanh = nn.Tanh()  # tanh激活函数
        self.embedding = nn.Embedding(self.num_classes, self.num_classes)  # embedding操作

    def forward(self, input: tuple):  # 前传函数
        noise, label = input  # 从输入的元组中获取噪声向量和标签信息
        label = self.embedding(label)  # 标签信息经过embedding操作
        z = torch.cat([noise, label], dim=-1)  # 噪声向量与标签向量进行拼接，(n,100)cat(n,2)-->(n,102)
        z = self.linear1(z)  # linear映射,(n,102)-->(n,256)
        z = self.leakyrelu(z)  # leakyrelu激活函数
        z = self.linear2(z)  # linear映射,(n,256)-->(n,512)
        z = self.bn2(z)  # 一维bn操作
        z = self.leakyrelu(z)  # leakyrelu激活函数
        z = self.linear3(z)  # linear映射,(n,512)-->(n,1024)
        z = self.bn3(z)  # 一维bn操作
        z = self.leakyrelu(z)  # leakyrelu激活函数
        z = self.linear4(z)  # linear映射,(n,1024)-->(n,784)
        z = self.tanh(z)  # tanh激活函数，将输出压缩至（-1.1）
        z = z.view(-1, 1, self.img_size[0], self.img_size[1])  # 从一维特征向量扩展至三维图片，(n,784)-->(n,1,28,28)

        return z  # 返回生成的图片