from jittor.dataset.mnist import MNIST import jittor.transform as transform from jittor.dataset.dataset import ImageFolder import jittor as jt from jittor import nn, Module import os import argparse from time import * import PIL.Image as Image import numpy as np import matplotlib.pyplot as plt plt.switch_backend('agg') jt.flags.use_cuda = 1 # 参数设定 parser = argparse.ArgumentParser() 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='训练数据集地址') parser.add_argument('--eval_dir', type=str, default='D:\\Image_Generation_Learn\\Dataset\\CelebA_train', help='训练数据集地址') parser.add_argument('--n_epochs', type=int, default=100, 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('--img_size', type=int, default=112, help='每个图像尺寸的大小') parser.add_argument('--celebA_channels', type=int, default=3, help='图像通道数') parser.add_argument('--mnist_channels', type=int, default=1, help='图像通道数') parser.add_argument('--img_row', type=int, default=5, help='图像样本之间的间隔') parser.add_argument('--img_column', type=int, default=5, 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('--sample_interval', type=int, default=400, help='图像样本之间的间隔') ''' opt = parser.parse_args() print(opt) # 训练集加载程序 def DataLoader(dataclass, img_size, batch_size, train_dir, eval_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) eval_loader = MNIST (data_root=eval_dir, train=False, transform = Transform).set_attrs(batch_size=1, 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) eval_loader = ImageFolder(eval_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, eval_dir) return train_loader, eval_loader # 加载训练集数据 train_loader, eval_loader = DataLoader(dataclass=opt.task,img_size=opt.img_size,batch_size=opt.batch_size,train_dir=opt.train_dir,eval_dir=opt.eval_dir) # 生成器 class generator(Module): def __init__(self, dim=3): super(generator, self).__init__() self.fc = nn.Linear(1024, 7*7*256) self.fc_bn = nn.BatchNorm(256) self.deconv1 = nn.ConvTranspose(256, 256, 3, 2, 1, 1) self.deconv1_bn = nn.BatchNorm(256) self.deconv2 = nn.ConvTranspose(256, 256, 3, 1, 1) self.deconv2_bn = nn.BatchNorm(256) self.deconv3 = nn.ConvTranspose(256, 256, 3, 2, 1, 1) self.deconv3_bn = nn.BatchNorm(256) self.deconv4 = nn.ConvTranspose(256, 256, 3, 1, 1) self.deconv4_bn = nn.BatchNorm(256) self.deconv5 = nn.ConvTranspose(256, 128, 3, 2, 1, 1) self.deconv5_bn = nn.BatchNorm(128) self.deconv6 = nn.ConvTranspose(128, 64, 3, 2, 1, 1) self.deconv6_bn = nn.BatchNorm(64) self.deconv7 = nn.ConvTranspose(64 , dim, 3, 1, 1) self.relu = nn.ReLU() self.tanh = nn.Tanh() def execute(self, input): x = self.fc(input).reshape((-1, 256, 7, 7)) x = self.relu(self.fc_bn(x)) x = self.relu(self.deconv1_bn(self.deconv1(x))) x = self.relu(self.deconv2_bn(self.deconv2(x))) x = self.relu(self.deconv3_bn(self.deconv3(x))) x = self.relu(self.deconv4_bn(self.deconv4(x))) x = self.relu(self.deconv5_bn(self.deconv5(x))) x = self.relu(self.deconv6_bn(self.deconv6(x))) x = self.tanh(self.deconv7(x)) return x # 判别器 class discriminator(nn.Module): def __init__(self, dim=3): super(discriminator, self).__init__() self.conv1 = nn.Conv(dim, 64, 5, 2, 2) self.conv2 = nn.Conv(64, 128, 5, 2, 2) self.conv2_bn = nn.BatchNorm(128) self.conv3 = nn.Conv(128, 256, 5, 2, 2) self.conv3_bn = nn.BatchNorm(256) self.conv4 = nn.Conv(256, 512, 5, 2, 2) self.conv4_bn = nn.BatchNorm(512) self.fc = nn.Linear(512*7*7, 1) self.leaky_relu = nn.Leaky_relu() def execute(self, input): x = self.leaky_relu(self.conv1(input), 0.2) x = self.leaky_relu(self.conv2_bn(self.conv2(x)), 0.2) x = self.leaky_relu(self.conv3_bn(self.conv3(x)), 0.2) x = self.leaky_relu(self.conv4_bn(self.conv4(x)), 0.2) x = x.reshape((x.shape[0], 512*7*7)) x = self.fc(x) return x # 损失函数 def ls_loss(x, b): mini_batch = x.shape[0] y_real_ = jt.ones((mini_batch,)) y_fake_ = jt.zeros((mini_batch,)) if b: return (x-y_real_).sqr().mean() else: return (x-y_fake_).sqr().mean() # 定义图像拼接函数 def image_compose(array,IMAGE_SIZE=128,IMAGE_SAVE_PATH='./images_celebA'): to_image = Image.new('RGB', (opt.img_column * IMAGE_SIZE, opt.img_row * IMAGE_SIZE)) # 创建一个新图 randomList = np.random.randint(0,array.shape[0],25) img_list = list() for i in randomList: # print(type(array[i])) img = Image.fromarray(np.uint8(array[i].transpose((1,2,0))*255)) img_list.append(img) # 循环遍历，把每张图片按顺序粘贴到对应位置上 for y in range(1, opt.img_row + 1): for x in range(1, opt.img_column + 1): from_image = img_list.pop().resize((IMAGE_SIZE, IMAGE_SIZE), Image.ANTIALIAS) to_image.paste(from_image, ((x - 1) * IMAGE_SIZE, (y - 1) * IMAGE_SIZE)) return to_image.save(IMAGE_SAVE_PATH) # 保存新图 def save_img_result(num_epoch, G, path = './images_celebA/result.png'): fixed_z_ = jt.init.gauss((5 * 5, 1024), 'float') # fixed noise z_ = fixed_z_ G.eval() test_images = G(z_) G.train() size_figure_grid = 5 fig, ax = plt.subplots(size_figure_grid, size_figure_grid, figsize=(5, 5)) for i in range(size_figure_grid): for j in range(size_figure_grid): ax[i, j].get_xaxis().set_visible(False) ax[i, j].get_yaxis().set_visible(False) for k in range(5*5): i = k // 5 j = k % 5 ax[i, j].cla() if opt.task=="MNIST": ax[i, j].imshow((test_images[k, 0].data+1)/2, cmap='gray') else: ax[i, j].imshow((test_images[k].data.transpose(1, 2, 0)+1)/2) label = 'Epoch {0}'.format(num_epoch) fig.text(0.5, 0.04, label, ha='center') plt.savefig(path) plt.close() def train(epoch): for batch_idx, (x_, target) in enumerate(train_loader): mini_batch = x_.shape[0] # 判别器训练 将假图片尽可能的判别为0 D_result = D(x_) #输入[128,3,112,112,] 生成[128,1] 128位batch_size D_real_loss = ls_loss(D_result, True) #真实图片的损失 z_ = jt.init.gauss((mini_batch, 1024), 'float') #生成随机噪声，大小为[128,1024] G_result = G(z_) #输入噪声，生成[128,3,112,112,] D_result_ = D(G_result) #输入由噪声生成的图像，得到判别器的预测值 D_fake_loss = ls_loss(D_result_, False) #假图片的损失 D_train_loss = D_real_loss + D_fake_loss D_train_loss.sync() D_optim.step(D_train_loss) # 生成器训练 让生成器尽可能的生成真实的照片 z_ = jt.init.gauss((mini_batch, 1024), 'float') #生成噪声 G_result = G(z_) #由噪声生成假图片 D_result = D(G_result) #将假图片输入到判别器，得到预测值 G_train_loss = ls_loss(D_result, True) #将假图片的预测值与1做损失，目的是未来让生成器尽可能的生成真实的照片 G_train_loss.sync() G_optim.step(G_train_loss) if (batch_idx%100==0 ): print("train: epoch{} batch_idx{} D training loss = {} G training loss = {} ".format(epoch,batch_idx,D_train_loss.data.mean(),G_train_loss.data.mean())) # if((epoch)%5==0 or epoch==0 and batch_idx==100): # image_compose(G_result.data,128,"./imgs/epoch{}-G_{}.jpg".format(epoch,task)) def validate(epoch): D_losses = [] G_losses = [] G.eval() D.eval() for batch_idx, (x_, target) in enumerate(eval_loader): mini_batch = x_.shape[0] # 判别器损失计算 D_result = D(x_) D_real_loss = ls_loss(D_result, True) z_ = jt.init.gauss((mini_batch, 1024), 'float') G_result = G(z_) D_result_ = D(G_result) D_fake_loss = ls_loss(D_result_, False) D_train_loss = D_real_loss + D_fake_loss D_losses.append(D_train_loss.data.mean()) # 生成器损失计算 z_ = jt.init.gauss((mini_batch, 1024), 'float') G_result = G(z_) D_result = D(G_result) G_train_loss = ls_loss(D_result, True) G_losses.append(G_train_loss.data.mean()) G.train() D.train() print("validate: epoch{}\tbatch_idx{}\tD training loss = {}\tG training loss = {}" .format(epoch, batch_idx, str(np.array(D_losses).mean()), str(np.array(G_losses).mean()))) # 初始化生成器和判别器 (通道数) G = generator(opt.celebA_channels) D = discriminator(opt.celebA_channels) # 优化器 0.0002 (0.5, 0.999) G_optim = jt.nn.Adam(G.parameters(), opt.lr, betas=(opt.b1, opt.b2)) D_optim = jt.nn.Adam(D.parameters(), opt.lr, betas=(opt.b1, opt.b2)) # 结果存储地址 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) G.load_parameters(jt.load(save_model_path+'/generator_celebA.pkl')) D.load_parameters(jt.load(save_model_path+'/discriminator_celebA.pkl')) for epoch in range(37,opt.n_epochs): print ('number of epochs', epoch) train(epoch) #validate(epoch) result_img_path = save_img_path + '/' + str(epoch) + '.png' save_img_result(epoch, G, path=result_img_path) # 指定地址保存训练好的模型 if (epoch+1) % 10 == 0: G.save(save_model_path+"/generator_celebA.pkl") D.save(save_model_path+"/discriminator_celebA.pkl")