Upload gan_model_load_test.py

gan_model_load_test.py

@@ -0,0 +1,101 @@
+import jittor as jt
+from jittor import init
+from jittor import nn
+
+import argparse
+import numpy as np
+import cv2
+
+jt.flags.use_cuda = 1
+
+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)
+
+# 生成器
+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
+
+def deal_image(img, path=None, 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))
+    if path:
+        # 根据地址保存图片样本数据
+        cv2.imwrite(path,img)
+    cv2.imshow('1',img)
+    cv2.waitKey(0)
+
+# 初始化生成器和判别器,并加载模型
+generator = Generator()
+g_model_path = "saved_models/generator_last.pkl"
+generator.load_parameters(jt.load(g_model_path))
+generator.load(g_model_path)
+discriminator = Discriminator()
+d_model_path = "saved_models/discriminator_last.pkl"
+discriminator.load_parameters(jt.load(d_model_path))
+discriminator.load(d_model_path)
+
+z = jt.array(np.random.normal(0, 1, (64, opt.latent_dim)).astype(np.float32))
+gen_imgs = generator(z)
+deal_image(gen_imgs.data[:25], "images_test/1.png",nrow=5)