File size: 7,034 Bytes
58da73e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import itertools
import torch
from util.image_pool import ImagePool
from . import networks
from .base_model import BaseModel
class CycleGANModel(BaseModel):
"""
This class implements the CycleGAN model, for learning image-to-image translation without paired data.
The model training requires '--dataset_mode unaligned' dataset.
By default, it uses a '--netG resnet_9blocks' ResNet generator,
a '--netD basic' discriminator (PatchGAN introduced by pix2pix),
and a least-square GANs objective ('--gan_mode lsgan').
CycleGAN paper: https://arxiv.org/pdf/1703.10593.pdf
"""
@staticmethod
def modify_commandline_options(parser, is_train=True):
"""Add new dataset-specific options, and rewrite default values for existing options.
Parameters:
parser -- original option parser
parser:
is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
Returns:
the modified parser.
For CycleGAN, in addition to GAN losses, we introduce lambda_A, lambda_B, and lambda_identity for the following losses.
A (source domain), B (target domain).
Generators: G_A: A -> B; G_B: B -> A.
Discriminators: D_A: G_A(A) vs. B; D_B: G_B(B) vs. A.
Forward cycle loss: lambda_A * ||G_B(G_A(A)) - A|| (Eqn. (2) in the paper)
Backward cycle loss: lambda_B * ||G_A(G_B(B)) - B|| (Eqn. (2) in the paper)
Identity loss (optional): lambda_identity * (||G_A(B) - B|| * lambda_B + ||G_B(A) - A|| * lambda_A) (Sec 5.2 "Photo generation from paintings" in the paper)
Dropout is not used in the original CycleGAN paper.
"""
parser.set_defaults(no_dropout=True) # default CycleGAN did not use dropout
if is_train:
parser.add_argument(
"--lambda_A",
type=float,
default=10.0,
help="weight for cycle loss (A -> B -> A)",
)
parser.add_argument(
"--lambda_B",
type=float,
default=10.0,
help="weight for cycle loss (B -> A -> B)",
)
parser.add_argument(
"--lambda_identity",
type=float,
default=0.5,
help="use identity mapping. Setting lambda_identity other than 0 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set lambda_identity = 0.1",
)
return parser
def __init__(self, opt):
"""Initialize the CycleGAN class.
Parameters:
opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions
"""
BaseModel.__init__(self, opt)
# specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals>
visual_names_A = ["real_A", "fake_B", "rec_A"]
visual_names_B = ["real_B", "fake_A", "rec_B"]
if (
self.isTrain and self.opt.lambda_identity > 0.0
): # if identity loss is used, we also visualize idt_B=G_A(B) ad idt_A=G_A(B)
visual_names_A.append("idt_B")
visual_names_B.append("idt_A")
self.visual_names = (
visual_names_A + visual_names_B
) # combine visualizations for A and B
# specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks>.
self.net_names = ["G_A", "G_B"]
if self.isTrain:
self.net_names.extend(["D_A", "D_B"])
# 下面会根据 self.loss_names self.visual_names net_names 中定义的字符串创建对应的变量名
# 这样把变量名写在一个列表中而不用字典，通过字典取得变量的写法可能是为了避免代码写的太长？
# 关键字 exec 可以根据字符串新建变量
# 用法：
# varlist = ["a"]
# exec
# define networks (both Generators and discriminators)
# The naming is different from those used in the paper.
# Code (vs. paper): G_A (G), G_B (F), D_A (D_Y), D_B (D_X)
self.net_G_A = networks.define_G(
opt.input_nc,
opt.output_nc,
opt.ngf,
opt.netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
opt.init_gain,
self.gpu_ids,
)
self.net_G_B = networks.define_G(
opt.output_nc,
opt.input_nc,
opt.ngf,
opt.netG,
opt.norm,
not opt.no_dropout,
opt.init_type,
opt.init_gain,
self.gpu_ids,
)
if self.isTrain: # define discriminators
self.net_D_A = networks.define_D(
opt.output_nc,
opt.ndf,
opt.netD,
opt.n_layers_D,
opt.norm,
opt.init_type,
opt.init_gain,
self.gpu_ids,
)
self.net_D_B = networks.define_D(
opt.input_nc,
opt.ndf,
opt.netD,
opt.n_layers_D,
opt.norm,
opt.init_type,
opt.init_gain,
self.gpu_ids,
)
if self.isTrain:
if (
opt.lambda_identity > 0.0
): # only works when input and output images have the same number of channels
assert opt.input_nc == opt.output_nc
self.fake_A_pool = ImagePool(
opt.pool_size
) # create image buffer to store previously generated images
self.fake_B_pool = ImagePool(
opt.pool_size
) # create image buffer to store previously generated images
def set_input(self, input):
"""Unpack input data from the dataloader and perform necessary pre-processing steps.
Parameters:
input (dict): include the data itself and its metadata information.
The option 'direction' can be used to swap domain A and domain B.
"""
AtoB = self.opt.direction == "AtoB"
self.real_A = input["A" if AtoB else "B"].to(self.device)
self.real_B = input["B" if AtoB else "A"].to(self.device)
self.image_paths = input["A_paths" if AtoB else "B_paths"]
def forward(self):
"""Run forward pass; called by both functions <optimize_parameters> and <test>."""
self.fake_B = self.net_G_A(self.real_A) # G_A(A)
self.rec_A = self.net_G_B(self.fake_B) # G_B(G_A(A))
self.fake_A = self.net_G_B(self.real_B) # G_B(B)
self.rec_B = self.net_G_A(self.fake_A) # G_A(G_B(B))
|