AIGC入门之生成对抗网络

"""  
Discriminator and Generator implementation from DCGAN paper  
"""  
  
import torch  
import torch.nn as nn  
from torchinfo import summary  
  
  
class Discriminator(nn.Module):  
    def __init__(self, channels_img, features_d):  
        super().__init__()  
        self.disc = nn.Sequential(  
            self._block(channels_img, features_d, kernel_size=4, stride=2, padding=1),  
            self._block(features_d, features_d * 2, 4, 2, 1),  
            self._block(features_d * 2, features_d * 4, 4, 2, 1),  
            self._block(features_d * 4, features_d * 8, 4, 2, 1),  
            nn.Conv2d(features_d * 8, 1, kernel_size=4, stride=2, padding=0),  
            nn.Sigmoid(),  
        )  
  
    def _block(self, in_channels, out_channels, kernel_size, stride, padding):  
        return nn.Sequential(  
            nn.Conv2d(  
                in_channels,  
                out_channels,  
                kernel_size,  
                stride,  
                padding,  
                bias=False  
            ),  
            nn.LeakyReLU(0.2),  
        )  
  
    def forward(self, x):  
        return self.disc(x)  
  
  
class Generator(nn.Module):  
    def __init__(self, channels_noise, channels_img, features_g):  
        super().__init__()  
        self.gen = nn.Sequential(  
            self._block(channels_noise, features_g * 16, 4, 1, 0),  
            self._block(features_g * 16, features_g * 8, 4, 2, 1),  
            self._block(features_g * 8, features_g * 4, 4, 2, 1),  
            self._block(features_g * 4, features_g * 2, 4, 2, 1),  
            nn.ConvTranspose2d(features_g * 2, channels_img, 4, 2, 1),  
            nn.Tanh(),  
        )  
  
    def _block(self, in_channels, out_channels, kernel_size, stride, padding):  
        return nn.Sequential(  
            nn.ConvTranspose2d(  
                in_channels,  
                out_channels,  
                kernel_size,  
                stride,  
                padding,  
                bias=False,  
            ),  
            nn.ReLU(),  
        )  
  
    def forward(self, x):  
        return self.gen(x)  
  
  
def initialize_weights(model):  
    ## initilialize weight according to paper  
    for m in model.modules():  
        if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d,)):  
            nn.init.normal_(m.weight.data, 0.0, 0.02)  
  
  
def test():  
    N, in_channels, H, W = 8, 1, 64, 64  
    noise_dim = 100  
    x = torch.randn(N, in_channels, H, W)  
    disc = Discriminator(in_channels, 8)  
    initialize_weights(disc)  
    assert disc(x).shape == (N, 1, 1, 1), "Discriminator test failed"  
    gen = Generator(noise_dim, in_channels, 8)  
    initialize_weights(gen)  
    z = torch.randn(N, noise_dim, 1, 1)  
    assert gen(z).shape == (N, in_channels, H, W), "Generator test failed"  
  
  
if __name__ == "__main__":  
    gnet = Generator(100, 1, 64)  
    dnet = Discriminator(1, 64)  
  
    summary(gnet, input_data=[torch.randn(10, 100, 1, 1)])  
    summary(dnet, input_data=[torch.randn(10, 1, 64, 64)])

import torch  
import torch.nn as nn  
import torch.optim as optim  
from torchvision.datasets import MNIST  
from torch.utils.data import DataLoader  
from dcgan import Generator, Discriminator, initialize_weights  
import torchvision.transforms as transforms  
import torchvision  
import cv2  
import numpy as np  
  
device = "cuda" if torch.cuda.is_available() else "cpu"  
LEARNING_RATE = 1e-4  
BATCH_SIZE = 128  
IMAGE_SIZE = 64  
NUM_EPOCHS = 5  
CHANNELS_IMG = 1  
NOISE_DIM = 100  
FEATURES_DISC = 64  
FEATURES_GEN = 64  
  
transforms = transforms.Compose(  
    [  
        transforms.Resize(IMAGE_SIZE),  
  
        transforms.ToTensor(),  
        transforms.Normalize(  
            [0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)]  
        ),  
    ]  
)  
  
  
def train(NUM_EPOCHS):  
  
    # 数据load  
    dataset = MNIST(root='./data', train=True, download=True, transform=transforms)  
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)  
  
    gen = Generator(NOISE_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)  
    disc = Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)  
    initialize_weights(gen)  
    initialize_weights(disc)  
  
    opt_gen = optim.Adam(gen.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))  
    opt_disc = optim.Adam(disc.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))  
    criterion = nn.BCELoss()  
  
    fixed_noise = torch.randn(32, NOISE_DIM, 1, 1).to(device)  
  
    gen.train()  
    disc.train()  
  
    for epoch in range(NUM_EPOCHS):  
        # 不需要目标的标签，无监督  
        for batch_id, (real, _) in enumerate(dataloader):  
            real = real.to(device)  
            noise = torch.randn(BATCH_SIZE, NOISE_DIM, 1, 1).to(device)  
            fake = gen(noise)  
  
            # Train Discriminator: max log(D(x)) + log(1 - D(G(z)))  
            disc_real = disc(real).reshape(-1)  
            loss_real = criterion(disc_real, torch.ones_like(disc_real))  
  
            # 训练判别器，生成器输出的值从计算图剥离出来  
            disc_fake = disc(fake.detach()).reshape(-1)  
            loss_fake = criterion(disc_fake, torch.zeros_like(disc_fake))  
  
            loss_disc = (loss_real + loss_fake) / 2  
  
            disc.zero_grad()  
            loss_disc.backward()  
            opt_disc.step()  
  
            # Train Generator: min log(1 - D(G(z))) <-> max log(D(G(z)),  
            output = disc(fake).reshape(-1)  
            loss_gen = criterion(output, torch.ones_like(output))  
  
            gen.zero_grad()  
            loss_gen.backward()  
            opt_gen.step()  
  
            if batch_id % 20 == 0:  
                print(  
                    f'Epoch [{epoch}/{NUM_EPOCHS}] Batch {batch_id}/{len(dataloader)} Loss D: {loss_disc.item()}, loss G: {loss_gen.item()}')  
  
                # 推理生成结果  
                with torch.no_grad():  
                    fake = gen(fixed_noise)  
                    img_grid_real = torchvision.utils.make_grid(real[:32], normalize=True)  
                    img_grid_fake = torchvision.utils.make_grid(fake[:32], normalize=True)  
  
                    img_grid_combined = torch.cat((img_grid_real, img_grid_fake), dim=2)  
                    img_grid_combined = img_grid_combined.permute(1, 2, 0).cpu().detach().numpy()  
                    img_grid_combined = (img_grid_combined * 255).astype(np.uint8)  
                    # 使用 cv2 显示图片  
                    cv2.imshow('Combined Image', img_grid_combined)  
                    cv2.waitKey(1)  
                    cv2.imwrite(f"ckpt/tmp_dc.jpg", img_grid_combined)  
  
                # 防止生成器崩溃,重新初始话模型  
                if loss_gen.item() > 5:  
                    gen = Generator(NOISE_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)  
                    disc = Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)  
                    initialize_weights(gen)  
                    initialize_weights(disc)  
                    opt_gen = optim.Adam(gen.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))  
                    opt_disc = optim.Adam(disc.parameters(), lr=LEARNING_RATE, betas=(0.5, 0.999))  
  
        if (epoch + 1) % 10 == 0:  
            torch.save(gen.state_dict(), f"ckpt/dc_gen_weights.pth")  
            torch.save(disc.state_dict(), f"ckpt/dc_disc_weights.pth")  
  
  
if __name__ == "__main__":  
    train(100)