wgan_gp.py

import jittor as jt
from jittor import nn
import argparse
import os
import numpy as np
from jittor.dataset.mnist import MNIST
import jittor.transform as transform
import cv2
import time
from jittor.dataset.dataset import ImageFolder

jt.flags.use_cuda = 1

save_img_path = './images_celebA'
save_model_path = './save_model_celebA'
os.makedirs(save_img_path, exist_ok=True)
os.makedirs(save_model_path, exist_ok=True)

parser = argparse.ArgumentParser()
parser.add_argument('--n_epochs', type=int, default=200, help='训练的时期数')
parser.add_argument('--batch_size', type=int, default=128, help='批次大小')
parser.add_argument('--lr', type=float, default=0.0002, help='学习率')
parser.add_argument('--b1', type=float, default=0.5, help='梯度的一阶动量衰减')
parser.add_argument('--b2', type=float, default=0.999, help='梯度的一阶动量衰减')
parser.add_argument('--n_cpu', type=int, default=8, help='批处理生成期间要使用的 cpu 线程数')
parser.add_argument('--latent_dim', type=int, default=100, help='潜在空间的维度')
parser.add_argument('--img_size', type=int, default=28, help='每个图像尺寸的大小')
parser.add_argument('--celebA_channels', type=int, default=3, help='图像通道数')
parser.add_argument('--mnist_channels', type=int, default=1, help='图像通道数')
parser.add_argument('--n_critic', type=int, default=5, help='每个迭代器的鉴别器训练步骤数')
parser.add_argument('--clip_value', type=float, default=0.01, help='光盘的上下剪辑值。 权重')
parser.add_argument('--sample_interval', type=int, default=400, help='图像样本之间的间隔')
parser.add_argument('--task', type=str, default='celebA', help='训练数据集类型')
parser.add_argument('--train_dir', type=str, default='D:\\Image_Generation_Learn\\Dataset\\CelebA_train', help='训练数据集地址')
opt = parser.parse_args()
print(opt)
img_shape = (opt.celebA_channels, opt.img_size, opt.img_size)

# 训练集加载程序
def DataLoader(dataclass, img_size, batch_size, train_dir):
    if dataclass == 'MNIST':
        Transform = transform.Compose([
            transform.Resize(size=img_size),
            transform.Gray(),
            transform.ImageNormalize(mean=[0.5], std=[0.5])])
        train_loader = MNIST (data_root=train_dir, train=True, transform=Transform).set_attrs(batch_size=batch_size, shuffle=True)
    elif dataclass == 'celebA':
        Transform = transform.Compose([
            transform.Resize(size=img_size),
            transform.ImageNormalize(mean=[0.5, 0.5, 0.5],std=[0.5, 0.5, 0.5])])
        train_loader = ImageFolder(train_dir)\
            .set_attrs(transform=Transform, batch_size=batch_size, shuffle=True)
    else:
        print("没有加载%s数据集的程序,请选择MNIST或者celebA!" % dataclass)
        dataclass = input("请输入:MNIST或者celebA:")
        DataLoader(dataclass, img_size, batch_size,train_dir)
    
    return train_loader

dataloader = DataLoader(opt.task,opt.img_size,opt.batch_size,opt.train_dir)

# 保存图片
def save_image(img, path, nrow=10):
    N,C,W,H = img.shape
    img2=img.reshape([-1,W*nrow*nrow,H])
    img=img2[:,:W*nrow,:]
    for i in range(1,nrow):
        img=np.concatenate([img,img2[:,W*nrow*i:W*nrow*(i+1),:]],axis=2)
    min_=img.min()
    max_=img.max()
    img=(img-min_)/(max_-min_)*255
    img=img.transpose((1,2,0))
    cv2.imwrite(path,img)

# 生成器
class Generator(nn.Module):

    def __init__(self):
        super(Generator, self).__init__()

        def block(in_feat, out_feat, normalize=True):
            layers = [nn.Linear(in_feat, out_feat)]
            if normalize:
                layers.append(nn.BatchNorm1d(out_feat, 0.8))
            layers.append(nn.LeakyReLU(0.2))
            return layers
        self.model = nn.Sequential(*block(opt.latent_dim, 128, normalize=False), *block(128, 256), *block(256, 512), *block(512, 1024), nn.Linear(1024, int(np.prod(img_shape))), nn.Tanh())

    def execute(self, z):
        img = self.model(z)
        img = img.view((img.shape[0], *img_shape))
        return img

# 判别器
class Discriminator(nn.Module):

    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(nn.Linear(int(np.prod(img_shape)), 512),
                                   nn.LeakyReLU(0.2),
                                   nn.Linear(512, 256),
                                   nn.LeakyReLU(0.2),
                                   nn.Linear(256, 1),
                                   )

    def execute(self, img):
        img_flat = img.reshape((img.shape[0], (- 1)))
        validity = self.model(img_flat)
        return validity

lambda_gp = 10

# 初始化生成器和判别器
generator = Generator()
discriminator = Discriminator()

# 优化器
optimizer_G = jt.optim.Adam(generator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
optimizer_D = jt.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))

# 损失函数（计算 WGAN GP 的梯度惩罚损失）
def compute_gradient_penalty(D, real_samples, fake_samples):
    alpha = jt.array(np.random.random((real_samples.shape[0], 1, 1, 1)).astype('float32'))
    interpolates = ((alpha * real_samples) + ((1 - alpha) * fake_samples))
    d_interpolates = D(interpolates)
    gradients = jt.grad(d_interpolates, interpolates)
    gradients = gradients.reshape((gradients.shape[0], (- 1)))
    gp =((jt.sqrt((gradients.sqr()).sum(1))-1).sqr()).mean()
    return gp

batches_done = 0
warmup_times = -1
run_times = 3000
total_time = 0.
cnt = 0

# ----------
#  训练
# ----------

for epoch in range(opt.n_epochs):# 200
    for i, (imgs, _) in enumerate(dataloader):
        real_imgs = jt.array(imgs).float32()

        # -----------------
        #  训练生成器
        # -----------------

        z = jt.array((np.random.normal(0, 1, (imgs.shape[0], opt.latent_dim))).astype('float32'))
        fake_imgs = generator(z)
        real_validity = discriminator(real_imgs)
        fake_validity = discriminator(fake_imgs)
        gradient_penalty = compute_gradient_penalty(discriminator, real_imgs, fake_imgs)
        d_loss = (- real_validity.mean() + fake_validity.mean() + lambda_gp * gradient_penalty)
        d_loss.sync()
        optimizer_D.step(d_loss)
        
        # ---------------------
        #  训练判别器
        # ---------------------

        if ((i % opt.n_critic) == 0):
            fake_img = generator(z)
            fake_validityg = discriminator(fake_img)
            g_loss = -fake_validityg.mean()
            g_loss.sync()
            optimizer_G.step(g_loss)
            
            if warmup_times==-1:
                print(('[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]' % (epoch, opt.n_epochs, i, len(dataloader), d_loss.data, g_loss.data)))
                #if ((batches_done % opt.sample_interval) == 0):
                if ( i == 1583 ):#根据opt.batch_size而变化，每批次保存一次
                    save_image(fake_imgs.data[:25], ('%s/%d.png' % (save_img_path, batches_done)), nrow=5)
                batches_done += opt.n_critic

        if warmup_times!=-1:
            jt.sync_all()
            cnt += 1
            print(cnt)
            if cnt == warmup_times:
                jt.sync_all(True)
                sta = time.time()
            if cnt > warmup_times + run_times:
                jt.sync_all(True)
                total_time = time.time() - sta
                print(f"run {run_times} iters cost {total_time} seconds, and avg {total_time / run_times} one iter.")
                exit(0)
    
    if epoch % 10 == 0:# 0-199
        generator.save("%s/generator_%s.pkl"%(save_model_path, opt.task))
        discriminator.save("%s/discriminator_%s.pkl"%(save_model_path, opt.task))