网络效果

image.png

网络结构

这部分主要参考：https://github.com/aladdinpersson/Machine-Learning-Collection/tree/master/ML/Pytorch/GANs/CycleGAN

Generator

  1. import torch
  2. import torch.nn as nn
  4. class ConvBlock(nn.Module):
  5.     def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):
  6.         super().__init__()
  7.         self.conv = nn.Sequential(
  8.             nn.Conv2d(in_channels, out_channels, padding_mode="reflect", **kwargs)
  9.             if down
  10.             else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),
  11.             nn.InstanceNorm2d(out_channels),
  12.             nn.ReLU(inplace=True) if use_act else nn.Identity()
  13.         )
  15.     def forward(self,x):
  16.         return self.conv(x)
  19. class ResidualBlock(nn.Module):
  20.     def __init__(self, channels):
  21.         super().__init__()
  22.         self.block = nn.Sequential(
  23.             ConvBlock(channels, channels, kernel_size=3, padding=1),
  24.             ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),
  25.         )
  27.     def forward(self,x):
  28.         return x + self.block(x)
  31. class Generator(nn.Module):
  32.     def __init__(self, img_channels, num_features=64, num_residuals=6):
  33.         super().__init__()
  34.         self.initial = nn.Sequential(
  35.             nn.Conv2d(img_channels, num_features, kernel_size=7, stride=1, padding=3, padding_mode="reflect"),
  36.             nn.InstanceNorm2d(num_features),
  37.             nn.ReLU(inplace=True)
  38.         )
  40.         self.down_blocks = nn.ModuleList(
  41.             [
  42.                 ConvBlock(num_features,num_features*2, kernel_size=3, stride=2, padding=1),
  43.                 ConvBlock(num_features*2, num_features*4, kernel_size=3, stride=2, padding=1)
  44.             ]
  45.         )
  47.         self.residual_blocks = nn.Sequential(
  48.             *[ResidualBlock(num_features*4) for _ in range(num_residuals)]
  49.         )
  51.         self.up_blocks = nn.ModuleList(
  52.             [
  53.                 ConvBlock(num_features*4, num_features*2, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),
  54.                 ConvBlock(num_features*2, num_features, down=False, kernel_size=3, stride=2, padding=1, output_padding=1)
  55.             ]
  56.         )
  58.         self.last = nn.Conv2d(num_features, img_channels, kernel_size=7, stride=1, padding=3, padding_mode="reflect")
  60.     def forward(self,x):
  61.         x = self.initial(x)
  62.         for layer in self.down_blocks:
  63.             x = layer(x)
  64.         x = self.residual_blocks(x)
  65.         for layer in self.up_blocks:
  66.             x = layer(x)
  67.         x = self.last(x)
  69.         return torch.tanh(x)
  71. def test():
  72.     img_channels = 3
  73.     img_size = 256
  74.     x = torch.randn((2,img_channels,img_size,256))
  75.     model = Generator(img_channels,num_residuals=6)
  76.     preds = model(x)
  77.     print(preds.shape)
  79. if __name__ == "__main__":
  80.     test()

Discriminator

  1. import torch
  2. import torch.nn as nn
  5. class Block(nn.Module):
  6.     def __init__(self, in_channels, out_channels, stride):
  7.         super().__init__()
  8.         self.conv = nn.Sequential(
  9.             nn.Conv2d(in_channels,out_channels,4,stride,1,bias=True,padding_mode='reflect'),
  10.             nn.InstanceNorm2d(out_channels),
  11.             nn.LeakyReLU(0.2)
  12.         )
  14.     def forward(self,x):
  15.         return self.conv(x)
  18. ### Input image size: 3x256x256
  19. class Discriminator(nn.Module):
  20.     def __init__(self, in_channels=3, features=[64,128,256,512]):
  21.         super().__init__()
  22.         self.initial = nn.Sequential(
  23.             nn.Conv2d(
  24.                 in_channels,
  25.                 features[0],
  26.                 kernel_size=4,
  27.                 stride=2,
  28.                 padding=1,
  29.                 padding_mode="reflect"
  30.             ),
  31.             nn.LeakyReLU(0.2),
  32.         )
  34.         layers = []
  35.         in_channels = features[0]
  36.         for feature in features[1:]:
  37.             layers.append(Block(in_channels,feature,stride=1 if feature==features[-1] else 2))
  38.             in_channels = feature
  39.         layers.append(nn.Conv2d(in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode='reflect'))
  41.         self.model = nn.Sequential(*layers)
  44.     def forward(self,x):
  45.         x = self.initial(x)
  46.         x = self.model(x)
  47.         return torch.sigmoid(x)
  51. def test():
  52.     x = torch.randn((5,3,256,256))
  53.     model = Discriminator()
  54.     preds = model(x)
  55.     print(preds.shape)
  57. if __name__ == "__main__":
  58.     test()

Dataset

  1. from PIL import Image
  2. import os
  3. from torch.utils.data import Dataset
  4. import numpy as np
  7. class HorseZebraDataset(Dataset):
  8.     def __init__(self, root_zebra, root_horse, transform=None):
  9.         self.root_zebra = root_zebra
  10.         self.root_horse = root_horse
  11.         self.transform = transform
  13.         self.zebra_images = os.listdir(root_zebra)
  14.         self.horse_images = os.listdir(root_horse)
  16.         self.zebra_len = len(self.zebra_images)
  17.         self.horse_len = len(self.horse_images)
  18.         self.length_dataset = max(self.zebra_len, self.horse_len)
  20.     def __len__(self):
  21.         return self.length_dataset
  23.     def __getitem__(self, idx):
  24.         zebra_img = self.zebra_images[idx % self.zebra_len]
  25.         horse_img = self.horse_images[idx % self.horse_len]
  27.         zebra_path = os.path.join(self.root_zebra, zebra_img)
  28.         horse_path = os.path.join(self.root_horse, horse_img)
  30.         zebra_img = np.array(Image.open(zebra_path).convert("RGB"))
  31.         horse_img = np.array(Image.open(horse_path).convert("RGB"))
  33.         if self.transform:
  34.             augmentations = self.transform(image=zebra_img, image0=horse_img)
  36.             horse_img = augmentations["image0"]
  37.             zebra_img = augmentations["image"]
  39.         return {'A':horse_img , 'B': zebra_img}

网络训练

这部分主要参考：https://github.com/aitorzip/PyTorch-CycleGAN

Training

  1. import argparse
  2. import itertools
  3. from torch.utils.data import DataLoader
  4. from torch.autograd import Variable
  5. import torch
  6. import torch.nn as nn
  7. import torch.optim as optim
  9. from generator_model import Generator
  10. from discriminator_model import Discriminator
  12. from dataset import HorseZebraDataset
  13. import albumentations as A
  14. from albumentations.pytorch import ToTensorV2
  15. from utils import ReplayBuffer
  17. def main(opt):
  18.     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  20.     ### Load the data
  21.     # image pre-processing
  22.     transforms = A.Compose(
  23.         [
  24.             A.Resize(width=256, height=256),
  25.             A.HorizontalFlip(p=0.5),
  26.             A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
  27.             ToTensorV2(),
  28.         ],
  29.         additional_targets={"image0": "image"},
  30.     )
  32.     datasets = HorseZebraDataset(root_horse=opt.data_root + "trainA", root_zebra=opt.data_root + "trainB",
  33.                                  transform=transforms)
  36.     loader = DataLoader(datasets, batch_size=opt.batch_size, shuffle=True, num_workers=4)
  38.     ### Building the Network
  39.     netG_A2B = Generator(opt.input_nc).to(device)
  40.     netG_B2A = Generator(opt.input_nc).to(device)
  42.     netD_A = Discriminator(opt.input_nc).to(device)
  43.     netD_B = Discriminator(opt.input_nc).to(device)
  45.     # Losses
  46.     criterion_GAN = nn.MSELoss()
  47.     criterion_cycle = nn.L1Loss()
  48.     criterion_identity = nn.L1Loss()
  50.     # Optimizers & LR schedulers
  51.     optimizer_G = optim.Adam(itertools.chain(netG_A2B.parameters(), netG_B2A.parameters()),
  52.                              lr=opt.lr, betas=(0.5, 0.999))
  54.     optimizer_D_A = optim.Adam(netD_A.parameters(), lr=opt.lr, betas=(0.5, 0.999))
  55.     optimizer_D_B = optim.Adam(netD_B.parameters(), lr=opt.lr, betas=(0.5, 0.999))
  57.     ### Inputs & targets memory allocation
  58.     Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.Tensor
  59.     input_A = Tensor(opt.batch_size, opt.input_nc, opt.size, opt.size)
  60.     input_B = Tensor(opt.batch_size, opt.input_nc, opt.size, opt.size)
  62.     fake_A_buffer = ReplayBuffer()
  63.     fake_B_buffer = ReplayBuffer()
  64.     # Loss plot
  67.     for epoch in range(opt.n_epochs):
  68.         for idx, batch in enumerate(loader):
  69.             # set model input
  70.             real_A = Variable(input_A.copy_(batch['A']))
  71.             real_B = Variable(input_B.copy_(batch['B']))
  74.             ### generate A2B and B2A ###
  75.             optimizer_G.zero_grad()
  77.             # Identity loss
  78.             # G_A2B(B) should equal B if real B is fed
  79.             same_B = netG_A2B(real_B)
  80.             loss_identity_B = criterion_identity(same_B, real_B)*5.0
  81.             # G_B2A(A) should equal A if real A is fed
  82.             same_A = netG_B2A(real_A)
  83.             loss_identity_A = criterion_identity(same_A, real_A)*5.0
  85.             # GAN loss
  86.             fake_B = netG_A2B(real_A)
  87.             pred_fake = netD_B(fake_B)
  88.             loss_GAN_A2B = criterion_GAN(pred_fake, torch.ones_like(pred_fake))
  90.             fake_A = netG_B2A(real_B)
  91.             pred_fake = netD_A(fake_A)
  92.             loss_GAN_B2A = criterion_GAN(pred_fake, torch.ones_like(pred_fake))
  95.             # Cycle loss
  96.             cycle_A = netG_B2A(fake_B)
  97.             cycle_B = netG_A2B(fake_A)
  99.             loss_cycle = criterion_cycle(cycle_A, real_A) + criterion_cycle(cycle_B, real_B)
  100.             loss_cycle *= 10.0
  102.             loss_G = loss_identity_A + loss_identity_B + loss_GAN_A2B + loss_GAN_B2A + loss_cycle
  103.             loss_G.backward()
  105.             optimizer_G.step()
  108.             ### Discriminator A ###
  109.             optimizer_D_A.zero_grad()
  110.             # Real loss
  111.             pred_real = netD_A(real_A)
  112.             loss_D_real = criterion_GAN(pred_real, torch.ones_like(pred_real))
  114.             # Fake loss
  115.             fake_A = fake_A_buffer.push_and_pop(fake_A)
  116.             pred_fake = netD_A(fake_A.detach())
  117.             loss_D_fake = criterion_GAN(pred_fake, torch.zeros_like(pred_real))
  119.             # Total loss
  120.             loss_D_A = (loss_D_real + loss_D_fake)*0.5
  121.             loss_D_A.backward()
  123.             optimizer_D_A.step()
  125.             ### Discriminator B ###
  126.             optimizer_D_B.zero_grad()
  129.             # Real loss
  130.             pred_real = netD_B(real_B)
  131.             loss_D_real = criterion_GAN(pred_real, torch.ones_like(pred_real))
  133.             # Fake loss
  134.             fake_B = fake_B_buffer.push_and_pop(fake_B)
  135.             pred_fake = netD_B(fake_B.detach())
  136.             loss_D_fake = criterion_GAN(pred_fake, torch.zeros_like(pred_real))
  138.             # Total loss
  139.             loss_D_B = (loss_D_real + loss_D_fake)*0.5
  140.             loss_D_B.backward()
  142.             optimizer_D_B.step()
  144.             if idx % 50 == 0:
  145.                 print(
  146.                     f"Epoch [{epoch}/{opt.n_epochs}] Batch {idx}/{len(loader)} \
  147.                     Loss G: {loss_G:.4f}, loss_cycle: {loss_cycle:.4f}, loss_D_A: {loss_D_A:.4f},"
  148.                 )
  151.         torch.save(netG_A2B.state_dict(),'./output/netG_A2B.pth')
  152.         torch.save(netG_B2A.state_dict(), './output/netG_B2A.pth')
  153.         torch.save(netD_A.state_dict(), './output/netD_A.pth')
  154.         torch.save(netD_B.state_dict(), './output/netD_B.pth')
  158. if __name__ == "__main__":
  160.     parser = argparse.ArgumentParser()
  162.     parser.add_argument('--n_epochs', type=int, default=200, help="number of epochs of training")
  163.     parser.add_argument('--batch_size', type=int, default=2, help="size of the batches")
  164.     parser.add_argument('--data_root', type=str, default='./data/horse2zebra/', help="root directory of the dataset")
  165.     parser.add_argument('--lr', type=float, default=0.0002, help='initial learning rate')
  166.     parser.add_argument('--size', type=int, default=256, help='size of data crop(squared assumed)')
  167.     parser.add_argument('--input_nc', type=int, default=3, help='number of channels of input data')
  168.     parser.add_argument('--output_nc', type=int, default=3, help='number of channels of output data')
  169.     opt = parser.parse_args()
  170.     print(opt)
  172.     main(opt)

Predict

import argparse
import sys
import os
from PIL import Image
import torchvision.transforms as transforms
from torchvision.utils import save_image
from torch.utils.data import DataLoader
from torch.autograd import Variable
import torch
import numpy as np
from generator_model import Generator
from dataset import HorseZebraDataset

parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=1, help='size of the batches')
parser.add_argument('--dataroot', type=str, default='datasets/horse2zebra/', help='root directory of the dataset')
parser.add_argument('--input_nc', type=int, default=3, help='number of channels of input data')
parser.add_argument('--output_nc', type=int, default=3, help='number of channels of output data')
parser.add_argument('--size', type=int, default=256, help='size of the data (squared assumed)')
parser.add_argument('--n_cpu', type=int, default=1, help='number of cpu threads to use during batch generation')
parser.add_argument('--generator_A2B', type=str, default='./output/netG_A2B.pth', help='A2B generator checkpoint file')
parser.add_argument('--generator_B2A', type=str, default='./output/netG_B2A.pth', help='B2A generator checkpoint file')
opt = parser.parse_args()
print(opt)


###### Definition of variables ######
# Networks
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

netG_A2B = Generator(opt.input_nc).to(device)
netG_B2A = Generator(opt.output_nc).to(device)


# Load state dicts
netG_A2B.load_state_dict(torch.load(opt.generator_A2B))
netG_B2A.load_state_dict(torch.load(opt.generator_B2A))

# Set model's test mode
netG_A2B.eval()
netG_B2A.eval()

# Inputs & targets memory allocation
horse_path = "D:/d2l/CycleGAN/data/horse2zebra/trainA/n02381460_36.jpg"
horse_img = np.array(Image.open(horse_path).convert("RGB"))

zebra_path = "D:/d2l/CycleGAN/data/horse2zebra/trainB/n02391049_77.jpg"
zebra_img = np.array(Image.open(zebra_path).convert("RGB"))

# transforms.ToTensor()
transform = transforms.Compose([
    transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ]
)


real_A = transform(horse_img).unsqueeze(0).cuda()
real_B = transform(zebra_img).unsqueeze(0).cuda()


fake_A = 0.5*(netG_B2A(real_B).data + 1.0)
fake_B = 0.5*(netG_A2B(real_A).data + 1.0)

out = fake_B.squeeze().cpu().numpy()
img_1 = np.transpose(out, (1,2,0))

out = fake_A.squeeze().cpu().numpy()
img_2 = np.transpose(out, (1,2,0))

import matplotlib.pyplot as plt
plt.subplot(221),plt.imshow(horse_img),plt.title("input image"),plt.axis("off")
plt.subplot(222),plt.imshow(img_1),plt.title("output image"),plt.axis("off")
plt.subplot(223),plt.imshow(zebra_img),plt.title("input image"),plt.axis("off")
plt.subplot(224),plt.imshow(img_2),plt.title("output image"),plt.axis("off")

训练好的网络

output.zip