CycleGAN的简单实现（pytorch）

        CycleGAN是于2017年发表在ICCV上的由GAN发展而来的一种无监督机器学习算法，是一种实现图像风格转换功能的GAN网络，在此之前存在着pix2pix实现图像风格转换，但pix2pix具有很大的局限性，主要是要求针对两种风格图像要对应出现，而现实中很难找到一些风格不同相同图像，也能难去拍摄获得，CycleGan实现了这个功能，在两种类型图像之间进行转换，而不需要对应关系。比如把照片转换为油画风格，或者把照片的橘子转换为苹果、马与斑马之间的转换等。

实现效果：

 

马转斑马

代码实现：

网络定义和训练代码

'''
Descripttion: 
version: 
Author: MAPLE
Date: 2022-06-12 23:23:54
LastEditors: MAPLE
LastEditTime: 2022-06-28 23:24:09
'''
import os
import torch
import random
import torch.nn as nn
import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torch.nn import init
from torch.optim import lr_scheduler
from tqdm import tqdm
from torchvision.utils import save_image
import torch.optim as optim
import torchvision.transforms as transforms

torch.cuda.is_available()

def seed_torch(seed=2018):

    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TRAIN_DIR = "data/horse2zebra"
VAL_DIR = "data/horse2zebra"

BATCH_SIZE = 1
LEARNING_RATE = 2e-4#学习率
LAMBDA_IDENTITY = 5  # identityloss
LAMBDA_CYCLE = 10  # 循环一致性损失
NUM_WORKERS = 2
LOAD_MODEL = True#加载模型参数
SAVE_MODEL = True#保存模型参数

#模型参数保存位置
CHECKPOINT_GEN_H = "genh.pth.tar"
CHECKPOINT_GEN_Z = "genz.pth.tar"
CHECKPOINT_CRITIC_H = "critich.pth.tar"
CHECKPOINT_CRITIC_Z = "criticz.pth.tar"

#学习率调度超参数
EPOCH_COUNT = 1
N_EPOCHS = 100
N_EPOCHS_DECAY = 100

transforms = transforms.Compose(
    [
        transforms.Resize(286, Image.BICUBIC),#重构
        transforms.RandomCrop(256),#随机裁剪
        transforms.RandomHorizontalFlip(p=0.5),#随机水平翻转
        transforms.ToTensor(),#转成tensor格式
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))#归一化
    ]
)

# 自定义参数初始化方式，用于多层网络初始化
def init_weights(net, init_type='normal', init_gain=0.02):
    """Initialize network weights.
        使用标准正态分布
    """
    def init_func(m):  # define the initialization function
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, init_gain)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
        elif classname.find('BatchNorm2d') != -1:
            init.normal_(m.weight.data, 1.0, init_gain)
            init.constant_(m.bias.data, 0.0)

    print('initialize network with %s' % init_type)
    net.to(DEVICE)
    net.apply(init_func)  # apply the initialization function <init_func>

class ImagePool():
    def __init__(self, pool_size):
        self.pool_size = pool_size
        if self.pool_size > 0:  # create an empty pool
            self.num_imgs = 0
            self.images = []

    def query(self, images):
        """从缓存区返回图片
        """
        if self.pool_size == 0:  # if the buffer size is 0, do nothing
            return images
        return_images = []
        for image in images:
            image = torch.unsqueeze(image.data, 0)
            if self.num_imgs < self.pool_size:   # if the buffer is not full; keep inserting current images to the buffer
                self.num_imgs = self.num_imgs + 1
                self.images.append(image)
                return_images.append(image)
            else:
                p = random.uniform(0, 1)
                if p > 0.5:  # 50%的概率返回以前生成的图像
                    random_id = random.randint(
                        0, self.pool_size - 1)  # randint is inclusive
                    tmp = self.images[random_id].clone()
                    self.images[random_id] = image  # 将新得到的图片存入缓存区
                    return_images.append(tmp)
                else:       # by another 50% chance, the buffer will return the current image
                    return_images.append(image)
        # collect all the images and return
        return_images = torch.cat(return_images, 0)
        return return_images

GLOBAL_SEED = 1
def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
set_seed(GLOBAL_SEED)

# 两个结构链接区域使用Residual block模块，默认是9个重复模块
class ResnetBlock(nn.Module):
    """Define a Resnet block"""
    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Initialize the Resnet block
        """
        super(ResnetBlock, self).__init__()
        self.conv_block = self.build_conv_block(
            dim, padding_type, norm_layer, use_dropout, use_bias)

    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):
        """Construct a convolutional block."""
        conv_block = []
        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]

        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p,
                                 bias=use_bias), norm_layer(dim), nn.ReLU(True)]
        #根据经验得，dropout在卷积中一般没啥用
        if use_dropout:
            conv_block += [nn.Dropout(0.5)]
        p = 0
        if padding_type == 'reflect':
            conv_block += [nn.ReflectionPad2d(1)]

        conv_block += [nn.Conv2d(dim, dim, kernel_size=3,
                                 padding=p, bias=use_bias), norm_layer(dim)]
        return nn.Sequential(*conv_block)

    def forward(self, x):
        """Forward function (with skip connections)"""
        out = x + self.conv_block(x)  # add skip connections
        return out

# 使用Residual block的生成器
class ResnetGenerator(nn.Module):

    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=9, padding_type='reflect'):
        """Construct a Resnet-based generator
        """
        super(ResnetGenerator, self).__init__()

        use_bias = norm_layer == nn.InstanceNorm2d

        model = [nn.ReflectionPad2d(3),
                nn.Conv2d(input_nc, ngf, kernel_size=7,padding=0, bias=use_bias),
                 norm_layer(ngf),
                 nn.ReLU(True)]

        n_downsampling = 2
        for i in range(n_downsampling):  # add downsampling layers
            mult = 2 ** i
            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),
                      norm_layer(ngf * mult * 2),
                        nn.ReLU(True)]

        mult = 2 ** n_downsampling
        for i in range(n_blocks):       # add ResNet blocks

            model += [ResnetBlock(ngf * mult, padding_type=padding_type,
                                  norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]

        for i in range(n_downsampling):  # add upsampling layers
            mult = 2 ** (n_downsampling - i)
            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),kernel_size=3, stride=2,padding=1, output_padding=1,bias=use_bias),
                                 norm_layer(int(ngf * mult / 2)),
                                nn.ReLU(True)]
        model += [nn.ReflectionPad2d(3)]
        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
        model += [nn.Tanh()]

        self.model = nn.Sequential(*model)
        init_weights(self.model)

    def forward(self, input):
        """Standard forward"""
        return self.model(input)

#马尔可夫判别器（PatchGAN）,由卷积层构成，最后输出一个n*n的预测矩阵
class NLayerDiscriminator(nn.Module):
    """Defines a PatchGAN discriminator"""

    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):
        """Construct a PatchGAN discriminator"""
        super(NLayerDiscriminator, self).__init__()
        use_bias = norm_layer == nn.InstanceNorm2d

        kw = 4
        padw = 1
        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw,stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
        nf_mult = 1
        nf_mult_prev = 1
        for n in range(1, n_layers):  # gradually increase the number of filters
            nf_mult_prev = nf_mult
            nf_mult = min(2 ** n, 8)
            sequence += [
                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,kernel_size=kw, stride=2, padding=padw, bias=use_bias),
                norm_layer(ndf * nf_mult),
                nn.LeakyReLU(0.2, True)
            ]

        nf_mult_prev = nf_mult
        nf_mult = min(2 ** n_layers, 8)
        sequence += [
            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult,kernel_size=kw, stride=1, padding=padw, bias=use_bias),
            norm_layer(ndf * nf_mult),
            nn.LeakyReLU(0.2, True)
        ]
        # output 1 channel prediction map
        sequence += [nn.Conv2d(ndf * nf_mult, 1,kernel_size=kw, stride=1, padding=padw)]
        self.model = nn.Sequential(*sequence)
        init_weights(self.model)

    def forward(self, input):
        """Standard forward."""
        return self.model(input)


# 学习率调度

def get_scheduler(optimizer):
    """Return a learning rate scheduler
        前100个epoch保持不变，后100个epoch线性衰减到0
    """
    def lambda_rule(epoch):
        lr_l = 1.0 - max(0, epoch + EPOCH_COUNT -N_EPOCHS) / float(N_EPOCHS_DECAY + 1)
        return lr_l
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)
    return scheduler

def train_fn(disc_H, disc_Z, gen_H, gen_Z, loader, opt_disc, opt_gen, l1, mse):
    fake_H_pool = ImagePool(50)
    fake_Z_pool = ImagePool(50)
    H_reals = 0
    H_fakes = 0
    Z_reals = 0
    Z_fakes = 0
    loop = tqdm(loader, leave=True)

    for idx, data in enumerate(loop):
        zebra = data['B'].to(DEVICE)
        horse = data['A'].to(DEVICE)

        # Train Discriminators H and Z
        fake_horse = gen_H(zebra)
        fake_horse_train = fake_H_pool.query(fake_horse)
        D_H_real = disc_H(horse)
        D_H_fake = disc_H(fake_horse_train.detach())
        H_reals += D_H_real.mean().item()
        H_fakes += D_H_fake.mean().item()
        D_H_real_loss = mse(D_H_real, torch.ones_like(D_H_real))
        D_H_fake_loss = mse(D_H_fake, torch.zeros_like(D_H_fake))
        D_H_loss = D_H_real_loss + D_H_fake_loss

        fake_zebra = gen_Z(horse)
        fake_zebra_train = fake_Z_pool.query(fake_zebra)
        D_Z_real = disc_Z(zebra)
        D_Z_fake = disc_Z(fake_zebra_train.detach())
        Z_reals += D_Z_real.mean().item()
        Z_fakes += D_Z_fake.mean().item()
        D_Z_real_loss = mse(D_Z_real, torch.ones_like(D_Z_real))
        D_Z_fake_loss = mse(D_Z_fake, torch.zeros_like(D_Z_fake))
        D_Z_loss = D_Z_real_loss + D_Z_fake_loss

        # put it togethor
        D_loss = (D_H_loss + D_Z_loss)/2

        opt_disc.zero_grad()
        D_loss.backward()
        opt_disc.step()

        # Train Generators H and Z
        # adversarial loss for both generators
        D_H_fake = disc_H(fake_horse)
        D_Z_fake = disc_Z(fake_zebra)
        loss_G_H = mse(D_H_fake, torch.ones_like(D_H_fake))
        loss_G_Z = mse(D_Z_fake, torch.ones_like(D_Z_fake))

        # cycle loss
        cycle_zebra = gen_Z(fake_horse)
        cycle_horse = gen_H(fake_zebra)
        cycle_zebra_loss = l1(zebra, cycle_zebra)
        cycle_horse_loss = l1(horse, cycle_horse)

        # identity loss (remove these for efficiency if you set lambda_identity=0)
        identity_zebra = gen_Z(zebra)
        identity_horse = gen_H(horse)
        identity_zebra_loss = l1(zebra, identity_zebra)
        identity_horse_loss = l1(horse, identity_horse)

        # add all togethor
        G_loss = (
            loss_G_Z
            + loss_G_H
            + cycle_zebra_loss * LAMBDA_CYCLE
            + cycle_horse_loss * LAMBDA_CYCLE
            + identity_horse_loss * LAMBDA_IDENTITY
            + identity_zebra_loss * LAMBDA_IDENTITY
        )

        opt_gen.zero_grad()
        G_loss.backward()
        opt_gen.step()

        if idx % 200 == 0:
            save_image(fake_horse*0.5+0.5, f"train_images/horse_{idx}.png")
            save_image(fake_zebra*0.5+0.5, f"train_images/zebra_{idx}.png")

        loop.set_postfix(H_real=H_reals/(idx+1), H_fake=H_fakes /
                         (idx+1), D_real=Z_reals/(idx+1), D_fake=Z_fakes/(idx+1))

class CombineDataset(Dataset):
    def __init__(self, root_A, root_B, transform):
        self.root_A = root_A
        self.root_B = root_B
        self.transform = transform

        self.A_paths = os.listdir(root_A)
        self.B_paths = os.listdir(root_B)
        self.length_dataset = max(len(self.A_paths), len(self.B_paths))
        self.A_len = len(self.A_paths)
        self.B_len = len(self.B_paths)

    def __len__(self):
        return self.length_dataset

    def __getitem__(self, index):
        A_path = self.A_paths[index % self.A_len]
        B_path = self.B_paths[index % self.B_len]

        A_img = Image.open(self.root_A+A_path).convert("RGB")
        B_img = Image.open(self.root_B+B_path).convert("RGB")

        A = self.transform(A_img)
        B = self.transform(B_img)

        return {'A': A, 'B': B}

def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"):
    print("=> Saving checkpoint")
    checkpoint = {
        "state_dict": model.state_dict(),
        "optimizer": optimizer.state_dict(),
    }
    torch.save(checkpoint, filename)


def load_checkpoint(checkpoint_file, model, optimizer, lr):
    print("=> Loading checkpoint")
    checkpoint = torch.load(checkpoint_file, map_location=DEVICE)
    model.load_state_dict(checkpoint["state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer"])

    # 修改学习率，使用当前的学习率
    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

dataset = CombineDataset(root_A=TRAIN_DIR+"/trainA/",
                         root_B=TRAIN_DIR+"/trainB/", transform=transforms)
data_loader = DataLoader(
    dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=NUM_WORKERS,
    pin_memory=True
)
dataset_size = len(data_loader)
print('The number of training images = %d' % dataset_size)

disc_H = NLayerDiscriminator(input_nc=3).to(DEVICE)
disc_Z = NLayerDiscriminator(input_nc=3).to(DEVICE)
gen_Z = ResnetGenerator(input_nc=3, output_nc=3).to(DEVICE)
gen_H = ResnetGenerator(input_nc=3, output_nc=3).to(DEVICE)

opt_disc = optim.Adam(
    list(disc_H.parameters()) + list(disc_Z.parameters()),
    lr=LEARNING_RATE,
    betas=(0.5, 0.999),
)

opt_gen = optim.Adam(
    list(gen_Z.parameters()) + list(gen_H.parameters()),
    lr=LEARNING_RATE,
    betas=(0.5, 0.999),
)

scheduler_disc = get_scheduler(opt_disc)
scheduler_gen = get_scheduler(opt_gen)
L1 = nn.L1Loss()
mse = nn.MSELoss()

if LOAD_MODEL:
    load_checkpoint(
        CHECKPOINT_GEN_H, gen_H, opt_gen, LEARNING_RATE,
    )
    load_checkpoint(
        CHECKPOINT_GEN_Z, gen_Z, opt_gen, LEARNING_RATE,
    )
    load_checkpoint(
        CHECKPOINT_CRITIC_H, disc_H, opt_disc, LEARNING_RATE,
    )
    load_checkpoint(
        CHECKPOINT_CRITIC_Z, disc_Z, opt_disc, LEARNING_RATE,
    )

for epoch in range(EPOCH_COUNT, N_EPOCHS+N_EPOCHS_DECAY+1):

    train_fn(disc_H, disc_Z, gen_H, gen_Z,
             data_loader, opt_disc, opt_gen, L1, mse)
    scheduler_disc.step()
    scheduler_gen.step()
    if SAVE_MODEL:
        save_checkpoint(gen_H, opt_gen, filename=CHECKPOINT_GEN_H)
        save_checkpoint(gen_Z, opt_gen, filename=CHECKPOINT_GEN_Z)
        save_checkpoint(disc_H, opt_disc, filename=CHECKPOINT_CRITIC_H)
        save_checkpoint(disc_Z, opt_disc, filename=CHECKPOINT_CRITIC_Z)

 完整工程训练参数数据集若需要请留言。