图像风格迁移模型-CycleGAN

作者： FutureSI
日期： 2021.03
摘要： 本案例实现了CycleGAN模型用于风格迁移。

一、CycleGAN介绍

CycleGAN，即循环生成对抗网络，是一种用于图片风格迁移的模型。原来的图片风格迁移模型通过在两组一一匹配的图片进行上训练，来学习输入图片组与输出图片组的特征映射关系，从而实现将输入图片的特征迁移到输出图片上，比如将A组图片的斑马的条纹外观特征迁移到B组普通马匹图片上。但是，训练所要求的两组一一对应训练集图片往往难以获得。CycleGAN通过给GAN网络添加循环一致性损失（consistency loss）的方法打破了训练集图片数据的一一对应限制。

二、框架导入设置

  1. # 解压 ai studio 数据集（首次执行后注释）
  2. !unzip -qa -d ~/data/data10040/ ~/data/data10040/horse2zebra.zip
  4. # 如果用wget自行下载数据集需要自行添加训练集列表文件
  5. # !wget https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/horse2zebra.zip
  6. # !unzip -qa -d /home/aistudio/data/data10040/ horse2zebra.zip
  1. import paddle
  2. from paddle.io import Dataset, DataLoader, IterableDataset
  3. import numpy as np
  4. import cv2
  5. import random
  6. import time
  7. import warnings
  8. import matplotlib.pyplot as plt
  9. %matplotlib inline
  11. warnings.filterwarnings("ignore", category=Warning) # 过滤报警信息
  13. BATCH_SIZE = 1
  14. DATA_DIR = '/home/aistudio/data/data10040/horse2zebra/' # 设置训练集数据地址

三、准备数据集

  1. from PIL import Image
  2. from paddle.vision.transforms import RandomCrop
  4. # 处理图片数据：随机裁切、调整图片数据形状、归一化数据
  5. def data_transform(img, output_size):
  6.     h, w, _ = img.shape
  7.     assert h == w and h >= output_size # check picture size
  8.     # random crop
  9.     rc = RandomCrop(224)
  10.     img = rc(img)
  11.     # normalize
  12.     img = img / 255. * 2. - 1.
  13.     # from [H,W,C] to [C,H,W]
  14.     img = np.transpose(img, (2, 0, 1))
  15.     # data type
  16.     img = img.astype('float32')
  17.     return img
  19. # 定义horse2zebra数据集对象
  20. class H2ZDateset(Dataset):
  21.     def __init__(self, data_dir):
  22.         super(H2ZDateset, self).__init__()
  23.         self.data_dir = data_dir
  24.         self.pic_list_a = np.loadtxt(data_dir+'trainA.txt', dtype=np.str)
  25.         np.random.shuffle(self.pic_list_a)
  26.         self.pic_list_b = np.loadtxt(data_dir+'trainB.txt', dtype=np.str)
  27.         np.random.shuffle(self.pic_list_b)
  28.         self.pic_list_lenth = min(int(self.pic_list_a.shape[0]), int(self.pic_list_b.shape[0]))
  30.     def __getitem__(self, idx):
  31.         img_dir_a = self.data_dir+self.pic_list_a[idx]
  32.         img_a = cv2.imread(img_dir_a)
  33.         img_a = cv2.cvtColor(img_a, cv2.COLOR_BGR2RGB)
  34.         img_a = data_transform(img_a, 224)
  35.         img_dir_b = self.data_dir+self.pic_list_b[idx]
  36.         img_b = cv2.imread(img_dir_b)
  37.         img_b = cv2.cvtColor(img_b, cv2.COLOR_BGR2RGB)
  38.         img_b = data_transform(img_b, 224)
  39.         return np.array([img_a, img_b])
  41.     def __len__(self):
  42.         return self.pic_list_lenth
  44. # 定义图片loader
  45. h2zdateset = H2ZDateset(DATA_DIR)
  46. loader = DataLoader(h2zdateset, shuffle=True, batch_size=BATCH_SIZE, drop_last=False, num_workers=0, use_shared_memory=False)
  47. data = next(loader())[0]
  48. data = np.transpose(data, (1, 0, 2, 3, 4))
  49. print("读取的数据形状：", data.shape)
  1. 读取的数据形状： [2, 1, 3, 224, 224]

四、模型组网

4.1 定义辅助功能函数

判别器负责区分图片的“真假”。输入的是训练集图片，判别器的输出越趋近于数值1（即判别此图片为真)；如果输入的是生成器生成的图片，判别器的输出越趋近于数值0（即判别此图片为假)。这样，生成器就可以根据判别器输出的变化而计算梯度以优化生成网络。

  1. from PIL import Image
  2. import os
  4. # 打开图片
  5. def open_pic(file_name='./data/data10040/horse2zebra/testA/n02381460_1300.jpg'):
  6.     img = Image.open(file_name).resize((256, 256), Image.BILINEAR)
  7.     img = (np.array(img).astype('float32') / 255.0 - 0.5) / 0.5
  8.     img = img.transpose((2, 0, 1))
  9.     img = img.reshape((-1, img.shape[0], img.shape[1], img.shape[2]))
  10.     return img
  12. # 存储图片
  13. def save_pics(pics, file_name='tmp', save_path='./output/pics/', save_root_path='./output/'):
  14.     if not os.path.exists(save_root_path):
  15.         os.makedirs(save_root_path)
  16.     if not os.path.exists(save_path):
  17.         os.makedirs(save_path)
  18.     for i in range(len(pics)):
  19.         pics[i] = pics[i][0]
  20.     pic = np.concatenate(tuple(pics), axis=2)
  21.     pic = pic.transpose((1,2,0))
  22.     pic = (pic + 1) / 2
  23.     # plt.imshow(pic)
  24.     pic = np.clip(pic * 256, 0, 255)
  25.     img = Image.fromarray(pic.astype('uint8')).convert('RGB')
  26.     img.save(save_path+file_name+'.jpg')
  28. # 显示图片
  29. def show_pics(pics):
  30.     print(pics[0].shape)
  31.     plt.figure(figsize=(3 * len(pics), 3), dpi=80)
  32.     for i in range(len(pics)):
  33.         pics[i] = (pics[i][0].transpose((1,2,0)) + 1) / 2
  34.         plt.subplot(1, len(pics), i + 1)
  35.         plt.imshow(pics[i])
  36.         plt.xticks([])
  37.         plt.yticks([])
  39. # 图片缓存队列
  40. class ImagePool(object):
  41.     def __init__(self, pool_size=50):
  42.         self.pool = []
  43.         self.count = 0
  44.         self.pool_size = pool_size
  45.     def pool_image(self, image):
  46.         return image
  47.         image = image.numpy()
  48.         rtn = ''
  49.         if self.count < self.pool_size:
  50.             self.pool.append(image)
  51.             self.count += 1
  52.             rtn = image
  53.         else:
  54.             p = np.random.rand()
  55.             if p > 0.5:
  56.                 random_id = np.random.randint(0, self.pool_size - 1)
  57.                 temp = self.pool[random_id]
  58.                 self.pool[random_id] = image
  59.                 rtn = temp
  60.             else:
  61.                 rtn = image
  62.         return paddle.to_tensor(rtn)

4.2 查看读取的数据集图片

  1. show_pics([data[0].numpy(), data[1].numpy()])
  1. (1, 3, 224, 224)

png

4.3 定义判别器

  1. import paddle
  2. import paddle.nn as nn
  3. import numpy as np
  5. # 定义基础的“卷积层+实例归一化”块
  6. class ConvIN(nn.Layer):
  7.     def __init__(self, num_channels, num_filters, filter_size, stride=1, padding=1, bias_attr=None,
  8.         weight_attr=None):
  9.         super(ConvIN, self).__init__()
  10.         model = [
  11.             nn.Conv2D(num_channels, num_filters, filter_size, stride=stride, padding=padding,
  12.                 bias_attr=bias_attr, weight_attr=weight_attr),
  13.             nn.InstanceNorm2D(num_filters),
  14.             nn.LeakyReLU(negative_slope=0.2)
  15.         ]
  16.         self.model = nn.Sequential(*model)
  17.     def forward(self, x):
  18.         return self.model(x)
  20. # 定义CycleGAN的判别器
  21. class Disc(nn.Layer):
  22.     def __init__(self, weight_attr=nn.initializer.Normal(0., 0.02)):
  23.         super(Disc, self).__init__()
  24.         model = [
  25.             ConvIN(3, 64, 4, stride=2, padding=1, bias_attr=True, weight_attr=weight_attr),
  26.             ConvIN(64, 128, 4, stride=2, padding=1, bias_attr=False, weight_attr=weight_attr),
  27.             ConvIN(128, 256, 4, stride=2, padding=1, bias_attr=False, weight_attr=weight_attr),
  28.             ConvIN(256, 512, 4, stride=1, padding=1, bias_attr=False, weight_attr=weight_attr),
  29.             nn.Conv2D(512, 1, 4, stride=1, padding=1, bias_attr=True, weight_attr=weight_attr)
  30.         ]
  31.         self.model = nn.Sequential(*model)
  32.     def forward(self, x):
  33.         return self.model(x)

4.4 测试判别器模块

  1. ci = ConvIN(3, 3, 3, weight_attr=nn.initializer.Normal(0., 0.02))
  2. logit = ci(paddle.to_tensor(data[0]))
  3. print('ConvIN块输出的特征图形状：', logit.shape)
  5. d = Disc()
  6. logit = d(paddle.to_tensor(data[0]))
  7. print('判别器输出的特征图形状：', logit.shape)
  1. ConvIN块输出的特征图形状： [1, 3, 224, 224]
  2. 判别器输出的特征图形状： [1, 1, 26, 26]

4.5 定义生成器

  1. # 定义基础的“转置卷积层+实例归一化”块
  2. class ConvTransIN(nn.Layer):
  3.     def __init__(self, num_channels, num_filters, filter_size, stride=1, padding='same', padding_mode='constant',
  4.         bias_attr=None, weight_attr=None):
  5.         super(ConvTransIN, self).__init__()
  6.         model = [
  7.             nn.Conv2DTranspose(num_channels, num_filters, filter_size, stride=stride, padding=padding,
  8.                 bias_attr=bias_attr, weight_attr=weight_attr),
  9.             nn.InstanceNorm2D(num_filters),
  10.             nn.LeakyReLU(negative_slope=0.2)
  11.         ]
  12.         self.model = nn.Sequential(*model)
  13.     def forward(self, x):
  14.         return self.model(x)
  16. # 定义残差块
  17. class Residual(nn.Layer):
  18.     def __init__(self, dim, bias_attr=None, weight_attr=None):
  19.         super(Residual, self).__init__()
  20.         model = [
  21.             nn.Conv2D(dim, dim, 3, stride=1, padding=1, padding_mode='reflect', bias_attr=bias_attr,
  22.                 weight_attr=weight_attr),
  23.             nn.InstanceNorm2D(dim),
  24.             nn.LeakyReLU(negative_slope=0.2),
  25.         ]
  26.         self.model = nn.Sequential(*model)
  27.     def forward(self, x):
  28.         return x + self.model(x)
  30. # 定义CycleGAN的生成器
  31. class Gen(nn.Layer):
  32.     def __init__(self, base_dim=64, residual_num=7, downup_layer=2, weight_attr=nn.initializer.Normal(0., 0.02)):
  33.         super(Gen, self).__init__()
  34.         model=[
  35.             nn.Conv2D(3, base_dim, 7, stride=1, padding=3, padding_mode='reflect', bias_attr=False,
  36.                 weight_attr=weight_attr),
  37.             nn.InstanceNorm2D(base_dim),
  38.             nn.LeakyReLU(negative_slope=0.2)
  39.         ]
  40.         # 下采样块（down sampling）
  41.         for i in range(downup_layer):
  42.             model += [
  43.                 ConvIN(base_dim * 2 ** i, base_dim * 2 ** (i + 1), 3, stride=2, padding=1, bias_attr=False,
  44.                     weight_attr=weight_attr),
  45.             ]
  46.         # 残差块（residual blocks）
  47.         for i in range(residual_num):
  48.             model += [
  49.                 Residual(base_dim * 2 ** downup_layer, True, weight_attr=nn.initializer.Normal(0., 0.02))
  50.             ]
  51.         # 上采样块（up sampling）
  52.         for i in range(downup_layer):
  53.             model += [
  54.                 ConvTransIN(base_dim * 2 ** (downup_layer - i), base_dim * 2 ** (downup_layer - i - 1), 3,
  55.                     stride=2, padding='same', padding_mode='constant', bias_attr=False, weight_attr=weight_attr),
  56.             ]
  57.         model += [
  58.             nn.Conv2D(base_dim, 3, 7, stride=1, padding=3, padding_mode='reflect', bias_attr=True,
  59.                 weight_attr=nn.initializer.Normal(0., 0.02)),
  60.             nn.Tanh()
  61.         ]
  62.         self.model = nn.Sequential(*model)
  63.     def forward(self, x):
  64.         return self.model(x)

4.6 测试生成器模块

  1. cti = ConvTransIN(3, 3, 3, stride=2, padding='same', padding_mode='constant', bias_attr=False,
  2.     weight_attr=nn.initializer.Normal(0., 0.02))
  3. logit = cti(paddle.to_tensor(data[0]))
  4. print('ConvTransIN块输出的特征图形状：', logit.shape)
  6. r = Residual(3, True, weight_attr=nn.initializer.Normal(0., 0.02))
  7. logit = r(paddle.to_tensor(data[0]))
  8. print('Residual块输出的特征图形状：', logit.shape)
  10. g = Gen()
  11. logit = g(paddle.to_tensor(data[0]))
  12. print('生成器输出的特征图形状：', logit.shape)
  1. ConvTransIN块输出的特征图形状： [1, 3, 448, 448]
  2. Residual块输出的特征图形状： [1, 3, 224, 224]
  3. 生成器输出的特征图形状： [1, 3, 224, 224]

五、训练CycleGAN网络

  1. # 模型训练函数
  2. def train(epoch_num=99999, adv_weight=1, cycle_weight=10, identity_weight=10, \
  3.           load_model=False, model_path='./model/', model_path_bkp='./model_bkp/', \
  4.           print_interval=1, max_step=5, model_bkp_interval=2000):
  6.     # 定义两对生成器、判别器对象
  7.     g_a = Gen()
  8.     g_b = Gen()
  9.     d_a = Disc()
  10.     d_b = Disc()
  12.     # 定义数据读取器
  13.     dataset = H2ZDateset(DATA_DIR)
  14.     reader_ab = DataLoader(dataset, shuffle=True, batch_size=BATCH_SIZE, drop_last=False, num_workers=2)
  16.     # 定义优化器
  17.     g_a_optimizer = paddle.optimizer.Adam(learning_rate=0.0002, beta1=0.5, beta2=0.999, parameters=g_a.parameters())
  18.     g_b_optimizer = paddle.optimizer.Adam(learning_rate=0.0002, beta1=0.5, beta2=0.999, parameters=g_b.parameters())
  19.     d_a_optimizer = paddle.optimizer.Adam(learning_rate=0.0002, beta1=0.5, beta2=0.999, parameters=d_a.parameters())
  20.     d_b_optimizer = paddle.optimizer.Adam(learning_rate=0.0002, beta1=0.5, beta2=0.999, parameters=d_b.parameters())
  22.     # 定义图片缓存队列
  23.     fa_pool, fb_pool = ImagePool(), ImagePool()
  25.     # 定义总迭代次数为0
  26.     total_step_num = np.array([0])
  28.     # 加载存储的模型
  29.     if load_model == True:
  30.         ga_para_dict = paddle.load(model_path+'gen_b2a.pdparams')
  31.         g_a.set_state_dict(ga_para_dict)
  33.         gb_para_dict = paddle.load(model_path+'gen_a2b.pdparams')
  34.         g_b.set_state_dict(gb_para_dict)
  36.         da_para_dict = paddle.load(model_path+'dis_ga.pdparams')
  37.         d_a.set_state_dict(da_para_dict)
  39.         db_para_dict = paddle.load(model_path+'dis_gb.pdparams')
  40.         d_b.set_state_dict(db_para_dict)
  42.         total_step_num = np.load('./model/total_step_num.npy')
  44.     # 定义本次训练开始时的迭代次数
  45.     step = total_step_num[0]
  47.     # 开始模型训练循环
  48.     print('Start time :', time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), 'start step:', step + 1)
  49.     for epoch in range(epoch_num):
  50.         for data_ab in reader_ab:
  51.             step += 1
  53.             # 设置模型为训练模式，针对bn、dropout等进行不同处理
  54.             g_a.train()
  55.             g_b.train()
  56.             d_a.train()
  57.             d_b.train()
  59.             # 得到A、B组图片数据
  60.             data_ab = np.transpose(data_ab[0], (1, 0, 2, 3, 4))
  61.             img_ra = paddle.to_tensor(data_ab[0])
  62.             img_rb = paddle.to_tensor(data_ab[1])
  64.             # 训练判别器DA
  65.             d_loss_ra = paddle.mean((d_a(img_ra.detach()) - 1) ** 2)
  66.             d_loss_fa = paddle.mean(d_a(fa_pool.pool_image(g_a(img_rb.detach()))) ** 2)
  67.             da_loss = (d_loss_ra + d_loss_fa) * 0.5
  68.             da_loss.backward() # 反向更新梯度
  69.             d_a_optimizer.step() # 更新模型权重
  70.             d_a_optimizer.clear_grad() # 清除梯度
  72.             # 训练判别器DB
  73.             d_loss_rb = paddle.mean((d_b(img_rb.detach()) - 1) ** 2)
  74.             d_loss_fb = paddle.mean(d_b(fb_pool.pool_image(g_b(img_ra.detach()))) ** 2)
  75.             db_loss = (d_loss_rb + d_loss_fb) * 0.5
  76.             db_loss.backward()
  77.             d_b_optimizer.step()
  78.             d_b_optimizer.clear_grad()
  80.             # 训练生成器GA
  81.             ga_gan_loss = paddle.mean((d_a(g_a(img_rb.detach())) - 1) ** 2)
  82.             ga_cyc_loss = paddle.mean(paddle.abs(img_rb.detach() - g_b(g_a(img_rb.detach()))))
  83.             ga_ide_loss = paddle.mean(paddle.abs(img_ra.detach() - g_a(img_ra.detach())))
  84.             ga_loss = ga_gan_loss * adv_weight + ga_cyc_loss * cycle_weight + ga_ide_loss * identity_weight
  85.             ga_loss.backward()
  86.             g_a_optimizer.step()
  87.             g_a_optimizer.clear_grad()
  89.             # 训练生成器GB
  90.             gb_gan_loss = paddle.mean((d_b(g_b(img_ra.detach())) - 1) ** 2)
  91.             gb_cyc_loss = paddle.mean(paddle.abs(img_ra.detach() - g_a(g_b(img_ra.detach()))))
  92.             gb_ide_loss = paddle.mean(paddle.abs(img_rb.detach() - g_b(img_rb.detach())))
  93.             gb_loss = gb_gan_loss * adv_weight + gb_cyc_loss * cycle_weight + gb_ide_loss * identity_weight
  94.             gb_loss.backward()
  95.             g_b_optimizer.step()
  96.             g_b_optimizer.clear_grad()
  98.             # 存储训练过程中生成的图片
  99.             if step in range(1, 101):
  100.                 pic_save_interval = 1
  101.             elif step in range(101, 1001):
  102.                 pic_save_interval = 10
  103.             elif step in range(1001, 10001):
  104.                 pic_save_interval = 100
  105.             else:
  106.                 pic_save_interval = 500
  107.             if step % pic_save_interval == 0:
  108.                 save_pics([img_ra.numpy(), g_b(img_ra).numpy(), g_a(g_b(img_ra)).numpy(), g_b(img_rb).numpy(), \
  109.                             img_rb.numpy(), g_a(img_rb).numpy(), g_b(g_a(img_rb)).numpy(), g_a(img_ra).numpy()], \
  110.                             str(step))
  111.                 test_pic = open_pic()
  112.                 test_pic_pp = paddle.to_tensor(test_pic)
  113.                 save_pics([test_pic, g_b(test_pic_pp).numpy()], str(step), save_path='./output/pics_test/')
  115.             # 打印训练过程中的loss值和生成的图片
  116.             if step % print_interval == 0:
  117.                 print([step], \
  118.                         'DA:', da_loss.numpy(), \
  119.                         'DB:', db_loss.numpy(), \
  120.                         'GA:', ga_loss.numpy(), \
  121.                         'GB:', gb_loss.numpy(), \
  122.                         time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
  123.                 show_pics([img_ra.numpy(), g_b(img_ra).numpy(), g_a(g_b(img_ra)).numpy(), g_b(img_rb).numpy()])
  124.                 show_pics([img_rb.numpy(), g_a(img_rb).numpy(), g_b(g_a(img_rb)).numpy(), g_a(img_ra).numpy()])
  126.             # 定期备份模型
  127.             if step % model_bkp_interval == 0:
  128.                 paddle.save(g_a.state_dict(), model_path_bkp+'gen_b2a.pdparams')
  129.                 paddle.save(g_b.state_dict(), model_path_bkp+'gen_a2b.pdparams')
  130.                 paddle.save(d_a.state_dict(), model_path_bkp+'dis_ga.pdparams')
  131.                 paddle.save(d_b.state_dict(), model_path_bkp+'dis_gb.pdparams')
  132.                 np.save(model_path_bkp+'total_step_num', np.array([step]))
  134.             # 完成训练时存储模型
  135.             if step >= max_step + total_step_num[0]:
  136.                 paddle.save(g_a.state_dict(), model_path+'gen_b2a.pdparams')
  137.                 paddle.save(g_b.state_dict(), model_path+'gen_a2b.pdparams')
  138.                 paddle.save(d_a.state_dict(), model_path+'dis_ga.pdparams')
  139.                 paddle.save(d_b.state_dict(), model_path+'dis_gb.pdparams')
  140.                 np.save(model_path+'total_step_num', np.array([step]))
  141.                 print('End time :', time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()), 'End Step:', step)
  142.                 return
  144. # 从头训练
  145. train()
  147. # 继续训练
  148. # train(print_interval=1, max_step=5, load_model=True)
  149. # train(print_interval=500, max_step=20000, load_model=True)
  1. Start time : 2021-03-10 11:36:45 start step: 1
  2. [1] DA: [1.5323195] DB: [2.9221125] GA: [13.066509] GB: [20.061096] 2021-03-10 11:36:46
  3. (1, 3, 224, 224)
  4. (1, 3, 224, 224)
  5. [2] DA: [3.431984] DB: [4.0848613] GA: [13.800614] GB: [12.840221] 2021-03-10 11:36:46
  6. (1, 3, 224, 224)
  7. (1, 3, 224, 224)
  8. [3] DA: [3.3024106] DB: [2.2502034] GA: [12.881987] GB: [12.331587] 2021-03-10 11:36:47
  9. (1, 3, 224, 224)
  10. (1, 3, 224, 224)
  11. [4] DA: [3.911097] DB: [1.5154138] GA: [12.64529] GB: [14.333654] 2021-03-10 11:36:47
  12. (1, 3, 224, 224)
  13. (1, 3, 224, 224)
  14. [5] DA: [1.9493798] DB: [1.8769395] GA: [14.874502] GB: [11.431137] 2021-03-10 11:36:48
  15. (1, 3, 224, 224)
  16. (1, 3, 224, 224)
  17. End time : 2021-03-10 11:36:48 End Step: 5

png

png

png

png

png

png

png

png

png

png

六、用训练好的模型进行预测

  1. def infer(img_path, model_path='./model/'):
  2.     # 定义生成器对象
  3.     g_b = Gen()
  5.     # 设置模型为训练模式，针对bn、dropout等进行不同处理
  6.     g_b.eval()
  8.     # 读取存储的模型
  9.     gb_para_dict = paddle.load(model_path+'gen_a2b.pdparams')
  10.     g_b.set_state_dict(gb_para_dict)
  12.     # 读取图片数据
  13.     img_a = cv2.imread(img_path)
  14.     img_a = cv2.cvtColor(img_a, cv2.COLOR_BGR2RGB)
  15.     img_a = data_transform(img_a, 224)
  16.     img_a = paddle.to_tensor(np.array([img_a]))
  18.     # 正向计算进行推理
  19.     img_b = g_b(img_a)
  21.     # 打印输出输入、输出图片
  22.     print(img_a.numpy().shape, img_a.numpy().dtype)
  23.     show_pics([img_a.numpy(), img_b.numpy()])
  25. infer('./data/data10040/horse2zebra/testA/n02381460_1300.jpg')
  1. (1, 3, 224, 224) float32

png