Upload gan.py

gan.py

@@ -0,0 +1,259 @@
+import jittor as jt
+from jittor import init
+from jittor import nn
+from jittor.dataset.mnist import MNIST
+import jittor.transform as transform
+import argparse
+import os
+import numpy as np
+import math
+import time
+import cv2
+
+jt.flags.use_cuda = 1
+os.makedirs('images', exist_ok=True)
+os.makedirs("saved_models", exist_ok=True)
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--n_epochs', type=int, default=200, help='训练的时期数')
+parser.add_argument('--batch_size', type=int, default=64, 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('--channels', type=int, default=1, help='图像通道数')
+parser.add_argument('--sample_interval', type=int, default=400, help='图像样本之间的间隔')
+
+opt = parser.parse_args()
+print(opt)
+img_shape = (opt.channels, opt.img_size, opt.img_size)
+
+# 保存生成器生成的图片样本数据
+def save_image(img, path, nrow=None):
+    N,C,W,H = img.shape# (25, 1, 28, 28)
+    '''
+    [-1,700,28] , img2的形状(1,700,28)
+    img[0][0][0] = img2[0][0]
+    img2:[
+            [1*28]
+            ......(一共700个)    
+                    ](1,700,28)
+    '''
+    img2=img.reshape([-1,W*nrow*nrow,H])
+    # [:,:28*5,:],img:(1,140,28)
+    img=img2[:,:W*nrow,:]
+    for i in range(1,nrow):#[1,5)
+        '''
+        img(1,140,28),img2(1,700,28)
+        img从(1,140,28)->(1,140,28+28)->...->(1,140,28+28+28+28)=(1,140,140)
+        np.concatenate把两个三维数组合并
+        '''        
+        img=np.concatenate([img,img2[:,W*nrow*i:W*nrow*(i+1),:]],axis=2)
+    # img中的数据大小从(-1，1)--(+1)-->(0,2)--(/2)-->(0,1)--(*255)-->(0,255)转换成了像素值
+    img=(img+1.0)/2.0*255
+    # (1,140,140)--->(140,140,1)
+    # (channels通道数,imagesize,imagesize)转化为(imagesize,imagesize,channels通道数)
+    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(scale=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(scale=0.2), nn.Linear(512, 256), nn.LeakyReLU(scale=0.2), nn.Linear(256, 1), nn.Sigmoid())
+
+    def execute(self, img):
+        img_flat = img.view((img.shape[0], (- 1)))
+        validity = self.model(img_flat)
+        return validity
+
+# bce loss分类器 （b这里指的是binary，所以用于二分类问题）
+'''
+源码：
+class BCELoss(Module):
+    def __init__(self, weight=None, size_average=True):
+        self.weight = weight
+        self.size_average = size_average
+    def execute(self, output, target):
+        return bce_loss(output, target, self.weight, self.size_average)
+
+# weight:表示对loss中每个元素的加权权值,默认为None
+# size_average:指定输出的格式，包括'mean','sum'
+# output:判别器对生成的数据的判别结果(64*1)
+# target:判别器对真实的数据的判别结果(64*1)
+def bce_loss(output, target, weight=None, size_average=True):
+    # jt.maximum(x,y):返回x和y的元素最大值
+    # 公式:损失值 = -权重*[ 理想结果*log(判别结果) + (1-理想结果)*log(1-判别结果) ]
+    loss = - ( 
+                target * jt.log(jt.maximum(output, 1e-20)) 
+                + 
+                (1 - target) * jt.log(jt.maximum(1 - output, 1e-20))
+            )
+    if weight is not None:
+        loss *= weight    
+    if size_average:
+        return loss.mean()# 求均值
+    else:
+        return loss.sum()# 求和
+'''
+# 对抗性损失函数
+adversarial_loss = nn.BCELoss()
+
+# 初始化生成器和判别器
+generator = Generator()
+discriminator = Discriminator()
+
+# 配置数据加载器
+transform = transform.Compose([
+    transform.Resize(size=opt.img_size),
+    transform.Gray(),
+    transform.ImageNormalize(mean=[0.5], std=[0.5]),
+])
+dataloader = MNIST(train=True, transform=transform).set_attrs(batch_size=opt.batch_size, shuffle=True)
+
+# 优化器
+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))
+
+warmup_times = -1
+run_times = 3000
+total_time = 0.
+cnt = 0
+
+# ----------
+#  训练
+# ----------
+
+for epoch in range(opt.n_epochs):#[1,200),200次迭代
+    for (i, (real_imgs, _)) in enumerate(dataloader):
+
+        '''
+        valid表示真,全为1,fake表示假,全为0
+        img.shape[0]:图像的垂直尺寸(高度)h
+        [ [1.0]...(一共h个)...[1.0] ]  64*1的数组
+        '''
+        valid = jt.ones([real_imgs.shape[0], 1]).stop_grad()
+        fake = jt.zeros([real_imgs.shape[0], 1]).stop_grad()
+
+        # ---------------------
+        #  训练生成器
+        # ---------------------  
+
+        # TODO 第一步：生成服从正态分布的噪音数据
+        '''
+        随机生成一个符合正态分布的噪声,numpy.random.normal(loc=0.0, scale=1.0, size=None)
+        loc:正态分布的均值,对应着这个分布的中心,0说明这一个以Y轴为对称轴的正态分布
+        scale:正态分布的标准差,对应分布的宽度,scale越大,正态分布的曲线越矮胖,scale越小,曲线越高瘦
+        shape:(图片的高度h,潜在空间的维度100) == (64,100) == z.shape
+        '''
+        z = jt.array(np.random.normal(0, 1, (real_imgs.shape[0], opt.latent_dim)).astype(np.float32))
+        # TODO 第二步：生成器加载噪音数据生成图片数据[64,1,28,28]
+        '''
+        gen_imgs的形状:(64,1,28,28), 64*1中每个元素都是28*28
+        [
+            [28*28]
+            ......   (一共64个28*28)
+                    ]
+        '''
+        gen_imgs = generator(z)
+        # TODO 第三步：根据生成数据的判别结果和真的数据（都是64*1）计算损失值
+        '''
+        把生成的图片数据放进判别器中,让判别器对其进行分类,计算出数据可能是真实数据的概率值(0-1之间的数)
+        valid当作是判别器分类的结果,全为1说明判别器认为这个数据来源于真实图片
+        adversarial_loss会调用bce_loss求损失值
+        因为我们需要使生成器生成的数据越来越像真实的数据,所以我们需要这两个数据越来越相似[discriminator(gen_imgs)和valid]
+        loss(x,y)=-w*[ylogx+(1-y)log(1-x)]
+        生成器理想条件下,discriminator(gen_imgs)=1,loss(1,1)=0
+        '''
+        g_loss = adversarial_loss(discriminator(gen_imgs), valid)
+        # TODO 第四步：反向传播，训练生成器的参数
+        optimizer_G.step(g_loss)
+        
+        # ---------------------
+        #  训练判别器
+        # ---------------------
+        
+        # TODO 第一步：根据训练集中的数据和真的数据计算损失值
+        '''
+        real_imgs:加载的训练集数据
+        把训练集数据放进判别器，得到判别器对训练集数据的判别结果，计算出数据可能是真实数据的概率值
+        valid当作是判别器分类的结果,全为1说明判别器认为这个数据来源于真实图片
+        因为我们需要使判别器把训练集数据判别为真实数据,所以我们需要使这两个数据越来越相似[discriminator(real_imgs), valid]
+        loss(x,y)=-w*[ylogx+(1-y)log(1-x)]
+        判别器理想条件下,discriminator(real_imgs)=1,loss(1,1)=0
+        '''
+        real_loss = adversarial_loss(discriminator(real_imgs), valid)#
+        # TODO 第二步：根据生成数据的判别结果和假的数据（都是64*1）计算损失值
+        '''
+        gen_imgs:生成器生成的图片数据
+        把生成的图片数据放进判别器中,让判别器对其进行分类,计算出数据可能是真实数据的概率值(0-1之间的数)
+        fake当作是判别器分类的结果,全为0说明判别器认为这个数据来源于生成的数据,而不是真实现实中的数据
+        调用bce_loss求损失值
+        因为我们需要使判别器能识别出机器生成的图片数据,所以我们需要使这两个数越来越相似[discriminator(gen_imgs), fake]
+        loss(x,y)=-w*[ylogx+(1-y)log(1-x)]
+        判别器理想条件下,discriminator(gen_imgs)=0,loss(0,0)=0
+        '''
+        fake_loss = adversarial_loss(discriminator(gen_imgs), fake)# 
+        # TODO 第三步：对这两个损失值求平均
+        d_loss = ((real_loss + fake_loss) / 2)
+        # TODO 第四步：反向传播，训练判别器的参数
+        optimizer_D.step(d_loss)
+
+        # ---------------------
+        #  打印训练进度，打印生成器和判别器的损失值
+        #  保存生成器生成的图片样本数据
+        # ---------------------
+
+        if warmup_times==-1:
+            '''
+            D loss:判别器的损失值,越小越好(0-1的数)
+            G loss:生成器的损失值,越小越好(0-1的数)
+            numpy():把Var数据类型的数据转换成array数据类型
+            '''
+            print(('[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]' % (epoch, opt.n_epochs, i, len(dataloader), d_loss.numpy()[0], g_loss.numpy()[0])))
+            # [0,200) * 938 + [0,938) = [0,199*938+937] = [0,187599]
+            batches_done = ((epoch * len(dataloader)) + i)
+            # opt.sample_interval = 400 , 187599 / 400 = 468
+            if ((batches_done % opt.sample_interval) == 0):
+                # gen_imgs.data[:25] -> (25, 1, 28, 28)
+                save_image(gen_imgs.data[:25], ('images/GAN_images/%d.png' % batches_done), nrow=5)
+        else:
+            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+1) % 10 == 0:
+        generator.save("saved_models/generator_last.pkl")
+        discriminator.save("saved_models/discriminator_last.pkl")