def __init__(self, config):
# Input shape
self.imsize = config["data"]["imsize"]
self.imchannels = config["data"]["imchannels"]
self.imshape = (self.imsize, self.imsize, self.imchannels)
self.ressize = config["train"]["res-size"]
self.resshape = (self.ressize, self.ressize,
config["train"]["res-filters"])
self.save_path = os.path.join(config["paths"]["save"],
config["run-title"])
# Configure data loader
self.dataloader = Dataloader(config)
self.dataloader.load()
# Configure models
models = Models(config)
# Calculate output shape of D (PatchGAN)
patch = int(self.imsize / 2**4)
self.disc_patch = (patch, patch, 1)
# Loss weights
self.lambda_cycle = config["train"][
"lambda-cycle-loss"] # Cycle-consistency loss
self.lambda_adv = 0.1 * self.lambda_cycle
self.lambda_id = 0.01 * self.lambda_cycle
self.lambda_res = 0.01 * self.lambda_cycle
optimizer = Adam(0.0002, 0.5)
# Build and compile the discriminators
self.d_A = models.build_discriminator(domain="A")
self.d_B = models.build_discriminator(domain="B")
self.d_A.compile(loss='mse', optimizer=optimizer, metrics=['accuracy'])
self.d_B.compile(loss='mse', optimizer=optimizer, metrics=['accuracy'])
# Freeze discriminators
self.d_A.trainable = False
self.d_B.trainable = False
# Build streams
self.d_c_stream = models.build_generator(domain="AB",
network="disentangler",
stream="C")
self.d_r_stream = models.build_generator(domain="AR",
network="disentangler",
stream="R")
self.e_c_stream = models.build_generator(domain="BA",
network="entangler",
stream="C")
# Build and compile disentangler
self.disentangler = self.build_disentangler()
self.disentangler.compile(
loss=['mse', 'mae', 'mae'],
loss_weights=[self.lambda_adv, self.lambda_cycle, self.lambda_id],
optimizer=optimizer)
# Build and compile entangler
self.entangler = self.build_entangler()
self.entangler.compile(loss=["mse", "mae", "mae", "mae", "mae", "mae"],
loss_weights=[
self.lambda_adv, self.lambda_cycle,
self.lambda_id, self.lambda_res,
self.lambda_res, self.lambda_res
],
optimizer=optimizer)
