G9、ACGAN理论与实战

# 初始化神经网络权重的函数
def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
        torch.nn.init.constant_(m.bias.data, 0.0)

# 生成器网络类
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()

        # 为类别标签创建嵌入层
        self.label_emb = nn.Embedding(opt.n_classes, opt.latent_dim)

        # 计算上采样前的初始大小
        self.init_size = opt.img_size // 4  # Initial size before upsampling

        # 第一层线性层
        self.l1 = nn.Sequential(nn.Linear(opt.latent_dim, 128 * self.init_size ** 2))

        # 卷积层块
        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 128, 3, stride=1, padding=1),
            nn.BatchNorm2d(128, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 64, 3, stride=1, padding=1),
            nn.BatchNorm2d(64, 0.8),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, opt.channels, 3, stride=1, padding=1),
            nn.Tanh(),
        )

    def forward(self, noise, labels):
        # 将标签嵌入到噪声中
        gen_input = torch.mul(self.label_emb(labels), noise)

        # 通过第一层线性层
        out = self.l1(gen_input)

        # 重新整形为合适的形状
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)

        # 通过卷积层块生成图像
        img = self.conv_blocks(out)
        return img

# 判别器网络类
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()

        # 定义判别器块的函数
        def discriminator_block(in_filters, out_filters, bn=True):
            """返回每个判别器块的层"""
            block = [nn.Conv2d(in_filters, out_filters, 3, 2, 1), nn.LeakyReLU(0.2, inplace=True), nn.Dropout2d(0.25)]
            if bn:
                block.append(nn.BatchNorm2d(out_filters, 0.8))
            return block

        # 判别器的卷积层块
        self.conv_blocks = nn.Sequential(
            *discriminator_block(opt.channels, 16, bn=False),
            *discriminator_block(16, 32),
            *discriminator_block(32, 64),
            *discriminator_block(64, 128),
        )

        # 下采样后图像的高度和宽度
        ds_size = opt.img_size // 2 ** 4

        # 输出层
        self.adv_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, 1), nn.Sigmoid())
        self.aux_layer = nn.Sequential(nn.Linear(128 * ds_size ** 2, opt.n_classes), nn.Softmax())

    def forward(self, img):
        out = self.conv_blocks(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        label = self.aux_layer(out)

        return validity, label

# 定义损失函数
adversarial_loss = torch.nn.BCELoss()  # 二元交叉熵损失，用于对抗训练
auxiliary_loss = torch.nn.CrossEntropyLoss()  # 交叉熵损失，用于辅助分类

# 初始化生成器和鉴别器
generator = Generator()  # 创建生成器实例
discriminator = Discriminator()  # 创建鉴别器实例

# 如果使用GPU，将模型和损失函数移至GPU上
if cuda:
    generator.cuda()
    discriminator.cuda()
    adversarial_loss.cuda()
    auxiliary_loss.cuda()

# 初始化模型权重
generator.apply(weights_init_normal)  # 初始化生成器的权重
discriminator.apply(weights_init_normal)  # 初始化鉴别器的权重

# 配置数据加载器
os.makedirs("./data/mnist_G7", exist_ok=True)  # 创建存储MNIST数据集的文件夹
dataloader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "./data/mnist_G7",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.Resize(opt.img_size), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])]
        ),
    ),
    batch_size=opt.batch_size,
    shuffle=True,
)

# 优化器
optimizer_G = torch.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))  # 生成器的优化器
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))  # 鉴别器的优化器

# 根据是否使用GPU选择数据类型
FloatTensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if cuda else torch.LongTensor

# 保存生成图像的函数
def sample_image(n_row, batches_done):
    """保存从0到n_classes的生成数字的图像网格"""
    # 采样噪声
    z = Variable(FloatTensor(np.random.normal(0, 1, (n_row ** 2, opt.latent_dim))))
    # 为n行生成标签从0到n_classes
    labels = np.array([num for _ in range(n_row) for num in range(n_row)])
    labels = Variable(LongTensor(labels))
    gen_imgs = generator(z, labels)
    save_image(gen_imgs.data, "images/%d.png" % batches_done, nrow=n_row, normalize=True)