def lapacian_pyramid(gauss_pyramid):
prev_layer = None
lap_pyramid = []
for layer in gauss_pyramid:
if prev_layer is not None:
lap_pyramid.append(prev_layer - upsample(layer))
prev_layer = layer
lap_pyramid.append(gauss_pyramid[-1])
return lap_pyramid
def compile(self, optimizer_factory, functions=None):
if functions is None:
functions = ['train', 'generate']
if functions == "all":
functions = ['train', 'generate', 'debug']
print(functions)
masks_idx = [T.tensor4("{}_mask_idx".format(name)) for name in self.names]
masks_bw = [self.bw_mask(m) for m in masks_idx]
x_real = T.tensor4("x_real")
gaussian_pyr = gaussian_pyramid(x_real, nb_layers=3)
lapacian_pyr = lapacian_pyramid(gaussian_pyr)
alfa_losses = self.alfa.losses(gaussian_pyr[2], gen_conditional=[masks_bw[0]])
alfa_image = self.alfa.g_out
alfa_image_up = upsample(alfa_image)
alfa_mask_loss = mask_loss(masks_idx[0], alfa_image)
alfa_g_loss = alfa_mask_loss + alfa_losses[0]
alfa_g_updates = optimizer_factory().get_updates(
self.alfa.G.params, self.alfa.G.constraints, alfa_g_loss)
alfa_d_loss = alfa_losses[1]
alfa_d_updates = optimizer_factory().get_updates(
self.alfa.D.params, self.alfa.D.constraints, alfa_d_loss)
gauss_16_up = upsample(gaussian_pyr[2])
bravo_gen_images = T.concatenate([alfa_image_up, gauss_16_up])
bravo_losses = self.bravo.losses(
lapacian_pyr[1],
gen_conditional=[masks_bw[1], alfa_image_up],
dis_conditional=[bravo_gen_images])
bravo_laplace = self.bravo.g_out
bravo_image = bravo_laplace + alfa_image_up
bravo_image_up = upsample(bravo_image)
bravo_mask_loss = mask_loss(masks_idx[1], bravo_image)
bravo_g_loss = bravo_mask_loss + bravo_losses[0]
bravo_g_updates = optimizer_factory().get_updates(
self.bravo.G.params, self.bravo.G.constraints, bravo_g_loss)
bravo_d_loss = bravo_losses[1]
bravo_d_updates = optimizer_factory().get_updates(
self.bravo.D.params, self.bravo.D.constraints, bravo_d_loss)
gauss_32_up = upsample(gaussian_pyr[1])
charlie_gen_images = T.concatenate([bravo_image_up, gauss_32_up])
charlie_losses = self.charlie.losses(lapacian_pyr[0],
gen_conditional=[masks_bw[2], bravo_image_up],
dis_conditional=[charlie_gen_images])
charlie_laplace = self.charlie.g_out
charlie_image = charlie_laplace + bravo_image_up
charlie_mask_loss = mask_loss(masks_idx[2], charlie_image)
charlie_g_loss = charlie_mask_loss + charlie_losses[0]
charlie_g_updates = optimizer_factory().get_updates(
self.charlie.G.params, self.charlie.G.constraints, charlie_g_loss)
charlie_d_loss = charlie_losses[1]
charlie_d_updates = optimizer_factory().get_updates(
self.charlie.D.params, self.charlie.D.constraints, charlie_d_loss)
if 'train' in functions:
self._train_fn = theano.function(
[x_real] + masks_idx,
alfa_losses + (alfa_mask_loss, alfa_g_loss) +
bravo_losses + (bravo_mask_loss, bravo_g_loss) +
charlie_losses + (charlie_mask_loss, charlie_g_loss),
updates=alfa_g_updates + alfa_d_updates + bravo_g_updates +
bravo_d_updates + charlie_g_updates + charlie_d_updates)
if 'generate' in functions:
self._generate = theano.function(
masks_bw,
[charlie_image]
)
if 'debug' in functions:
self.debug_labels = ["mask_bw_16", "mask_bw_32", "mask_bw_64",
"gauss_64", "gauss_32", "gauss_16", "gauss_8",
"gauss_32_up", "gauss_16_up",
"laplace_64", "laplace_32",
"charlie_laplace", "bravo_laplace",
"charlie_img", "bravo_img", "alfa_img",
"bravo_img_up", "alfa_img_up",
"charlie_d", "bravo_d", "alfa_d"]
self._debug = theano.function(
[x_real] + masks_idx,
masks_bw + gaussian_pyr +
[gauss_32_up, gauss_16_up] +
lapacian_pyr[:2] +
[charlie_laplace, bravo_laplace] +
[charlie_image, bravo_image, alfa_image] +
[bravo_image_up, alfa_image_up] +
[self.charlie.d_out, self.bravo.d_out, self.alfa.d_out]
)
