目录
一、前言
二、数据集
三、网络结构
四、代码
（一）net
（二）train
（三）test
五、结果
（一）loss
（二）训练可视化
（三）测试结果
六、完整代码
一、前言
pix2pix对训练样本要求较高，需要成对的数据集，而这种样本的获取往往需要耗费很大精力。CycleGAN恰巧解决了该问题，实现两个domain之间的转换，即只需要准备两种风格的数据集，让GAN去学习将domain X中的图片转换成domain Y的风格(不改变domain X原图中物体，仅仅实现风格转换)。
一种直观的思路是直接让G去学习domain X 到domain Y 以及domain Y 到domain X的映射关系，但这种方式会造成G生成图片的随机性太强，会使得生成的图片与输入的图片完全不相关，不仅违背了CycleGAN的目的，同时输出的结果也没有任何意义。
作者认为这种转换应该具有循环一致性，比如在语言翻译中，把一段话从中文翻译成英文，再从英文翻译回中文，意思应该是相近的，CycleGAN就是采用了这种思想。假设Ga表示Domain X到Domain Y的生成器，Gb表示Domain Y 到Domain X 的生成器，那么让Domain X中的图片real_A通过Ga后生成的图片fake_A再通过Gb生成的rec_A应该和A是高度相似的，Domain Y到Domain X同理。
CycleGAN中有两个生成器以及两个判别器，分别对应Domain X 到Domain Y 以及Domain Y到Domain X。
二、数据集
这里我采用的是monet2photo数据集(莫奈画->真实风景照片)，部分数据如下图所示。
Domain X(monet):
Domain Y(photo):
三、网络结构
生成器G的结构如下图所示，判别器D与pix2pix相同，网络结构pix2pix。
四、代码
（一）net
初始化方式与源码不同。
import torch.nn as nn
from torchsummary import summary
from collections import OrderedDict
# 定义残差块
class Resnet_block(nn.Module):
def __init__(self, in_channels):
super(Resnet_block, self).__init__()
block = []
for i in range(2):
block += [nn.ReflectionPad2d(1),
nn.Conv2d(in_channels, in_channels, 3, 1, 0),
nn.InstanceNorm2d(in_channels),
nn.ReLU(True) if i > 0 else nn.Identity()]
self.block = nn.Sequential(*block)
def forward(self, x):
out = x + self.block(x)
return out
class Cycle_Gan_G(nn.Module):
def __init__(self):
super(Cycle_Gan_G, self).__init__()
net_dic = OrderedDict()
# 三层卷积层
net_dic.update({'first layer': nn.Sequential(
nn.ReflectionPad2d(3), # [3,256,256] -> [3,262,262]
nn.Conv2d(3, 64, 7, 1), # [3,262,262] ->[64,256,256]
nn.InstanceNorm2d(64),
nn.ReLU(True)
)})
net_dic.update({'second_conv': nn.Sequential(
nn.Conv2d(64, 128, 3, 2, 1), # [128,128,128]
nn.InstanceNorm2d(128),
nn.ReLU(True)
)})
net_dic.update({'three_conv': nn.Sequential(
nn.Conv2d(128, 256, 3, 2, 1), # [256,64,64]
nn.InstanceNorm2d(256),
nn.ReLU(True)
)})
# 9层 resnet block
for i in range(6):
net_dic.update({'Resnet_block{}'.format(i + 1): Resnet_block(256)})
# up_sample
net_dic.update({'up_sample1': nn.Sequential(
nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.InstanceNorm2d(128), # [128,128,128]
nn.ReLU(True)
)})
net_dic.update({'up_sample2': nn.Sequential(
nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1),
nn.InstanceNorm2d(64), # [64,256,256]
nn.ReLU(True)
)})
net_dic.update({'last_layer': nn.Sequential(
nn.ReflectionPad2d(3),
nn.Conv2d(64, 3, 7, 1),
nn.Tanh()
)})
self.net_G = nn.Sequential(net_dic)
self.init_weight()
def init_weight(self):
for w in self.modules():
if isinstance(w, nn.Conv2d):
nn.init.kaiming_normal_(w.weight, mode='fan_out')
if w.bias is not None:
nn.init.zeros_(w.bias)
elif isinstance(w, nn.ConvTranspose2d):
nn.init.kaiming_normal_(w.weight, mode='fan_in')
elif isinstance(w, nn.BatchNorm2d):
nn.init.ones_(w.weight)
nn.init.zeros_(w.bias)
def forward(self, x):
out = self.net_G(x)
return out
class Cycle_Gan_D(nn.Module):
def __init__(self):
super(Cycle_Gan_D, self).__init__()
# 定义基本的卷积\bn\relu
def base_Conv_bn_lkrl(in_channels, out_channels, stride):
if in_channels == 3:
bn = nn.Identity
else:
bn = nn.InstanceNorm2d
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, 4, stride, 1),
bn(out_channels),
nn.LeakyReLU(0.2, True)
)
D_dic = OrderedDict()
in_channels = 3
out_channels = 64
for i in range(4):
if i < 3:
D_dic.update({'layer_{}'.format(i + 1): base_Conv_bn_lkrl(in_channels, out_channels, 2)})
else:
D_dic.update({'layer_{}'.format(i + 1): base_Conv_bn_lkrl(in_channels, out_channels, 1)})
in_channels = out_channels
out_channels *= 2
D_dic.update({'last_layer': nn.Conv2d(512, 1, 4, 1, 1)}) # [batch,1,30,30]
self.D_model = nn.Sequential(D_dic)
def forward(self, x):
return self.D_model(x)
if __name__ == '__main__':
# G = Cycle_Gan_G().to('cuda')
# summary(G, (3, 256, 256))
D = Cycle_Gan_D().to('cuda')
summary(D, (3, 256, 256))
（二）train
训练过程中有一些小细节，为了减小模型振荡，提高训练的稳定性，论文中采用了buffer来暂存G生成的图片，用之前生成的图片来更新判别器。G共包含三种损失(两个方向共6部分)，GAN_loss、Cycle_loss、id_loss。其中，GAN_loss就是传统GAN的loss，使得输出图片尽可能真，Cycle_loss是重建的图片与原始图片之间的L1损失，id_loss是为了保证G不去随意改变图片的色调(即便判别器告诉你另外一种色调也服从Domain Y的分布，但为了仅仅改变风格不改变别的因素，因此引入了该损失)。判别器D仍然采用了PatchGAN，训练过程与pix2pix类似。
import itertools
from image_pool import ImagePool
from torch.utils.tensorboard import SummaryWriter
from cyclegan import Cycle_Gan_G, Cycle_Gan_D
import argparse
from mydatasets import CreateDatasets
import os
from torch.utils.data.dataloader import DataLoader
import torch
import torch.optim as optim
import torch.nn as nn
from utils import train_one_epoch, val
def train(opt):
batch = opt.batch
data_path = opt.dataPath
print_every = opt.every
device = 'cuda' if torch.cuda.is_available() else 'cpu'
epochs = opt.epoch
img_size = opt.imgsize
if not os.path.exists(opt.savePath):
os.mkdir(opt.savePath)
# 加载数据集
train_datasets = CreateDatasets(data_path, img_size, mode='train')
val_datasets = CreateDatasets(data_path, img_size, mode='test')
train_loader = DataLoader(dataset=train_datasets, batch_size=batch, shuffle=True, num_workers=opt.numworker,
drop_last=True)
val_loader = DataLoader(dataset=val_datasets, batch_size=batch, shuffle=True, num_workers=opt.numworker,
drop_last=True)
# 实例化网络
Cycle_G_A = Cycle_Gan_G().to(device)
Cycle_D_A = Cycle_Gan_D().to(device)
Cycle_G_B = Cycle_Gan_G().to(device)
Cycle_D_B = Cycle_Gan_D().to(device)
# 定义优化器和损失函数
optim_G = optim.Adam(itertools.chain(Cycle_G_A.parameters(), Cycle_G_B.parameters()), lr=0.0002, betas=(0.5, 0.999))
optim_D = optim.Adam(itertools.chain(Cycle_D_A.parameters(), Cycle_D_B.parameters()), lr=0.0002, betas=(0.5, 0.999))
loss = nn.MSELoss()
l1_loss = nn.L1Loss()
start_epoch = 0
A_fake_pool = ImagePool(50)
B_fake_pool = ImagePool(50)
# 加载预训练权重
if opt.weight != '':
ckpt = torch.load(opt.weight)
Cycle_G_A.load_state_dict(ckpt['Ga_model'], strict=False)
Cycle_G_B.load_state_dict(ckpt['Gb_model'], strict=False)
Cycle_D_A.load_state_dict(ckpt['Da_model'], strict=False)
Cycle_D_B.load_state_dict(ckpt['Db_model'], strict=False)
start_epoch = ckpt['epoch'] + 1
writer = SummaryWriter('train_logs')
# 开始训练
for epoch in range(start_epoch, epochs):
loss_mG, loss_mD = train_one_epoch(Ga=Cycle_G_A, Da=Cycle_D_A, Gb=Cycle_G_B, Db=Cycle_D_B,
train_loader=train_loader,
optim_G=optim_G, optim_D=optim_D, writer=writer, loss=loss, device=device,
plot_every=print_every, epoch=epoch, l1_loss=l1_loss,
A_fake_pool=A_fake_pool, B_fake_pool=B_fake_pool)
writer.add_scalars(main_tag='train_loss', tag_scalar_dict={
'loss_G': loss_mG,
'loss_D': loss_mD
}, global_step=epoch)
# 保存模型
torch.save({
'Ga_model': Cycle_G_A.state_dict(),
'Gb_model': Cycle_G_B.state_dict(),
'Da_model': Cycle_D_A.state_dict(),
'Db_model': Cycle_D_B.state_dict(),
'epoch': epoch
}, './weights/cycle_monent2photo.pth')
# 验证集
val(Ga=Cycle_G_A, Da=Cycle_D_A, Gb=Cycle_G_B, Db=Cycle_D_B, val_loader=val_loader, loss=loss, l1_loss=l1_loss,
device=device, epoch=epoch)
def cfg():
parse = argparse.ArgumentParser()
parse.add_argument('--batch', type=int, default=1)
parse.add_argument('--epoch', type=int, default=100)
parse.add_argument('--imgsize', type=int, default=256)
parse.add_argument('--dataPath', type=str, default='../monet2photo', help='data root path')
parse.add_argument('--weight', type=str, default='', help='load pre train weight')
parse.add_argument('--savePath', type=str, default='./weights', help='weight save path')
parse.add_argument('--numworker', type=int, default=4)
parse.add_argument('--every', type=int, default=20, help='plot train result every * iters')
opt = parse.parse_args()
return opt
if __name__ == '__main__':
opt = cfg()
print(opt)
train(opt)
import torchvision
from tqdm import tqdm
import torch
import os
def train_one_epoch(Ga, Da, Gb, Db, train_loader, optim_G, optim_D, writer, loss, device, plot_every, epoch, l1_loss,
A_fake_pool, B_fake_pool):
pd = tqdm(train_loader)
loss_D, loss_G = 0, 0
step = 0
Ga.train()
Da.train()
Gb.train()
Db.train()
for idx, data in enumerate(pd):
A_real = data[0].to(device)
B_real = data[1].to(device)
# 前向传递
B_fake = Ga(A_real) # Ga生成的假B
A_rec = Gb(B_fake) # Gb重构回的A
A_fake = Gb(B_real) # Gb生成的假A
B_rec = Ga(A_fake) # Ga重构回的B
# 训练G => G包含六部分损失
set_required_grad([Da, Db], requires_grad=False) # 不更新D
optim_G.zero_grad()
ls_G = train_G(Da=Da, Db=Db, B_fake=B_fake, loss=loss, A_fake=A_fake, l1_loss=l1_loss,
A_rec=A_rec,
A_real=A_real, B_rec=B_rec, B_real=B_real, Ga=Ga, Gb=Gb)
ls_G.backward()
optim_G.step()
# 训练D
set_required_grad([Da, Db], requires_grad=True)
optim_D.zero_grad()
A_fake_p = A_fake_pool.query(A_fake)
B_fake_p = B_fake_pool.query(B_fake)
ls_D = train_D(Da=Da, Db=Db, B_fake=B_fake_p, B_real=B_real, loss=loss, A_fake=A_fake_p, A_real=A_real)
ls_D.backward()
optim_D.step()
loss_D += ls_D
loss_G += ls_G
pd.desc = 'train_{} G_loss: {} D_loss: {}'.format(epoch, ls_G.item(), ls_D.item())
# 绘制训练结果
if idx % plot_every == 0:
writer.add_images(tag='epoch{}_Ga'.format(epoch), img_tensor=0.5 * (torch.cat([A_real, B_fake], 0) + 1),
global_step=step)
writer.add_images(tag='epoch{}_Gb'.format(epoch), img_tensor=0.5 * (torch.cat([B_real, A_fake], 0) + 1),
global_step=step)
step += 1
mean_lsG = loss_G / len(train_loader)
mean_lsD = loss_D / len(train_loader)
return mean_lsG, mean_lsD
@torch.no_grad()
def val(Ga, Da, Gb, Db, val_loader, loss, device, l1_loss, epoch):
pd = tqdm(val_loader)
loss_D, loss_G = 0, 0
Ga.eval()
Da.eval()
Gb.eval()
Db.eval()
all_loss = 10000
for idx, item in enumerate(pd):
A_real_img = item[0].to(device)
B_real_img = item[1].to(device)
B_fake_img = Ga(A_real_img)
A_fake_img = Gb(B_real_img)
A_rec = Gb(B_fake_img)
B_rec = Ga(A_fake_img)
# D的loss
ls_D = train_D(Da=Da, Db=Db, B_fake=B_fake_img, B_real=B_real_img, loss=loss, A_fake=A_fake_img,
A_real=A_real_img)
# G的loss
ls_G = train_G(Da=Da, Db=Db, B_fake=B_fake_img, loss=loss, A_fake=A_fake_img, l1_loss=l1_loss,
A_rec=A_rec,
A_real=A_real_img, B_rec=B_rec, B_real=B_real_img, Ga=Ga, Gb=Gb)
loss_G += ls_G
loss_D += ls_D
pd.desc = 'val_{}: G_loss:{} D_Loss:{}'.format(epoch, ls_G.item(), ls_D.item())
# 保存最好的结果
all_ls = ls_G + ls_D
if all_ls < all_loss:
all_loss = all_ls
best_image = torch.cat([A_real_img, B_fake_img, B_real_img, A_fake_img], 0)
result_img = (best_image + 1) * 0.5
if not os.path.exists('./results'):
os.mkdir('./results')
torchvision.utils.save_image(result_img, './results/val_epoch{}_cycle.jpg'.format(epoch))
def set_required_grad(nets, requires_grad=False):
if not isinstance(nets, list):
nets = [nets]
for net in nets:
if net is not None:
for params in net.parameters():
params.requires_grad = requires_grad
def train_G(Da, Db, B_fake, loss, A_fake, l1_loss, A_rec, A_real, B_rec, B_real, Ga, Gb):
# GAN loss
Da_out_fake = Da(B_fake)
Ga_gan_loss = loss(Da_out_fake, torch.ones(Da_out_fake.size()).cuda())
Db_out_fake = Db(A_fake)
Gb_gan_loss = loss(Db_out_fake, torch.ones(Db_out_fake.size()).cuda())
# Cycle loss
Cycle_A_loss = l1_loss(A_rec, A_real) * 10
Cycle_B_loss = l1_loss(B_rec, B_real) * 10
# identity loss
Ga_id_out = Ga(B_real)
Gb_id_out = Gb(A_real)
Ga_id_loss = l1_loss(Ga_id_out, B_real) * 10 * 0.5
Gb_id_loss = l1_loss(Gb_id_out, A_real) * 10 * 0.5
# G的总损失
ls_G = Ga_gan_loss + Gb_gan_loss + Cycle_A_loss + Cycle_B_loss + Ga_id_loss + Gb_id_loss
return ls_G
def train_D(Da, Db, B_fake, B_real, loss, A_fake, A_real):
# Da的loss
Da_fake_out = Da(B_fake.detach()).squeeze()
Da_real_out = Da(B_real).squeeze()
ls_Da1 = loss(Da_fake_out, torch.zeros(Da_fake_out.size()).cuda())
ls_Da2 = loss(Da_real_out, torch.ones(Da_real_out.size()).cuda())
ls_Da = (ls_Da1 + ls_Da2) * 0.5
# Db的loss
Db_fake_out = Db(A_fake.detach()).squeeze()
Db_real_out = Db(A_real.detach()).squeeze()
ls_Db1 = loss(Db_fake_out, torch.zeros(Db_fake_out.size()).cuda())
ls_Db2 = loss(Db_real_out, torch.ones(Db_real_out.size()).cuda())
ls_Db = (ls_Db1 + ls_Db2) * 0.5
# D的总损失
ls_D = ls_Da + ls_Db
return ls_D
（三）test
from cyclegan import Cycle_Gan_G
import torch
import torchvision.transforms as transform
import matplotlib.pyplot as plt
import cv2
from PIL import Image
def test(img_path):
if img_path.endswith('.png'):
img = cv2.imread(img_path)
img = img[:, :, ::-1]
else:
img = Image.open(img_path)
transforms = transform.Compose([
transform.ToTensor(),
transform.Resize((256, 256)),
transform.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
img = transforms(img.copy())
img = img[None].to('cuda') # [1,3,128,128]
# 实例化网络
Gb = Cycle_Gan_G().to('cuda')
# 加载预训练权重
ckpt = torch.load('weights/cycle_monent2photo.pth')
Gb.load_state_dict(ckpt['Gb_model'], strict=False)
Gb.eval()
out = Gb(img)[0]
out = out.permute(1, 2, 0)
out = (0.5 * (out + 1)).cpu().detach().numpy()
plt.figure()
plt.imshow(out)
plt.show()
if __name__ == '__main__':
test('123.jpg')
五、结果
（一）loss
（二）训练可视化
这里我挑选了一部分训练结果和验证结果。
训练集上monet -> photo
训练集上photo-> monet
验证集上结果(左边为monet -> photo，右边为photo-> monet )
（三）测试结果
下图为photo转monet的结果
六、完整代码
数据集：百度网盘 请输入提取码 提取码：s3e3文章来源：https://www.toymoban.com/news/detail-442905.html
代码：百度网盘 请输入提取码 提取码：t0d5文章来源地址https://www.toymoban.com/news/detail-442905.html
到了这里，关于PyTorch 实现CycleGAN 风格迁移的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
