trainer.py

import os
import time
import tensorflow as tf
import numpy as np
from tqdm import trange
from math import ceil, log
from numpy import sin, cos
from scipy.misc import imsave


NO_REDUCTION = tf.losses.Reduction.NONE

class Trainer():
    def __init__(self, sess, config, real_images, 
                 g_builder, d_builder, cp_builder, zp_builder, 
                 coord_handler, patch_handler):
        self.sess = sess
        self.config = config
        self.real_images = real_images
        self.g_builder = g_builder
        self.d_builder = d_builder
        self.cp_builder = cp_builder
        self.zp_builder = zp_builder
        self.coord_handler = coord_handler
        self.patch_handler = patch_handler

        # Vars for graph building
        self.batch_size = self.config["train_params"]["batch_size"]
        self.z_dim = self.config["model_params"]["z_dim"]
        self.spatial_dim = self.config["model_params"]["spatial_dim"]
        self.micro_patch_size = self.config["data_params"]["micro_patch_size"]
        self.macro_patch_size = self.config["data_params"]["macro_patch_size"]

        self.ratio_macro_to_micro = self.config["data_params"]["ratio_macro_to_micro"]
        self.ratio_full_to_micro = self.config["data_params"]["ratio_full_to_micro"]
        self.num_micro_compose_macro = self.config["data_params"]["num_micro_compose_macro"]

        # Vars for training loop
        self.exp_name = config["log_params"]["exp_name"]
        self.epochs = float(self.config["train_params"]["epochs"])
        self.num_batches = self.config["data_params"]["num_train_samples"] // self.batch_size
        self.coordinate_system = self.config["data_params"]["coordinate_system"]
        self.G_update_period = self.config["train_params"]["G_update_period"]
        self.D_update_period = self.config["train_params"]["D_update_period"]
        self.Q_update_period = self.config["train_params"]["Q_update_period"]

        # Loss weights
        self.code_loss_w = self.config["loss_params"]["code_loss_w"]
        self.coord_loss_w = self.config["loss_params"]["coord_loss_w"]
        self.gp_lambda = self.config["loss_params"]["gp_lambda"]

        # Extrapolation parameters handling
        self.train_extrap = self.config["train_params"]["train_extrap"]
        if self.train_extrap:
            assert self.config["train_params"]["num_extrap_steps"] is not None
            assert self.coordinate_system is not "euclidean", \
                "I didn't handle extrapolation in {} coordinate system!".format(self.coordinate_system)
            self.num_extrap_steps = self.config["train_params"]["num_extrap_steps"]
        else:
            self.num_extrap_steps = 0


    def _train_content_prediction_model(self):
        return (self.Q_update_period>0) and (self.config["train_params"]["qlr"]>0)


    def sample_prior(self):
        return np.random.uniform(-1., 1., [self.batch_size, self.z_dim]).astype(np.float32)

    
    def _dup_z_for_macro(self, z):
        # Duplicate with nearest neighbor, different to `tf.tile`.
        # E.g., 
        # tensor: [[1, 2], [3, 4]]
        # repeat: 3
        # output: [[1, 2], [1, 2], [1, 2], [3, 4], [3, 4], [3, 4]]
        ch = z.shape[-1]
        repeat = self.num_micro_compose_macro
        extend = tf.expand_dims(z, 1)
        extend_dup = tf.tile(extend, [1, repeat, 1])
        return tf.reshape(extend_dup, [-1, ch])


    def build_graph(self, test_mode=False):

        # Input nodes
        # Note: the input node name was wrong in the checkpoint 
        self.micro_coord_fake = tf.placeholder(tf.float32, [None, self.spatial_dim], name='micro_coord_fake')
        self.macro_coord_fake = tf.placeholder(tf.float32, [None, self.spatial_dim], name='macro_coord_fake')
        self.micro_coord_real = tf.placeholder(tf.float32, [None, self.spatial_dim], name='micro_coord_real')
        self.macro_coord_real = tf.placeholder(tf.float32, [None, self.spatial_dim], name='macro_coord_real')

        # Reversing angle for cylindrical coordinate is complicated, directly pass values here
        self.y_angle_ratio = tf.placeholder(tf.float32, [None, 1], name='y_angle_ratio') 
        self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')
        
        # Crop real micro for visualization
        if self.coordinate_system == "euclidean":
             self.real_micro = self.patch_handler.crop_micro_from_full_gpu(
                self.real_images, self.micro_coord_real[:, 0:1], self.micro_coord_real[:, 1:2])
        elif self.coordinate_system == "cylindrical":
            self.real_micro = self.patch_handler.crop_micro_from_full_gpu(
                self.real_images, self.micro_coord_real[:, 0:1], self.y_angle_ratio)

        # Real part
        self.real_macro = self.patch_handler.concat_micro_patches_gpu(
            self.real_micro, ratio_over_micro=self.ratio_macro_to_micro)
        (self.disc_real, disc_real_h) = self.d_builder(self.real_macro, self.macro_coord_real, is_training=True)
        self.c_real_pred = self.cp_builder(disc_real_h, is_training=True)
        self.z_real_pred = self.zp_builder(disc_real_h, is_training=True)

        # Fake part
        z_dup_macro = self._dup_z_for_macro(self.z)
        self.gen_micro = self.g_builder(z_dup_macro, self.micro_coord_fake, is_training=True)
        self.gen_macro = self.patch_handler.concat_micro_patches_gpu(
            self.gen_micro, ratio_over_micro=self.ratio_macro_to_micro)
        (self.disc_fake, disc_fake_h) = self.d_builder(self.gen_macro, self.macro_coord_fake, is_training=True)
        self.c_fake_pred = self.cp_builder(disc_fake_h, is_training=True)
        self.z_fake_pred = self.zp_builder(disc_fake_h, is_training=True)

        # Testing graph
        if self.config["log_params"]["merge_micro_patches_in_cpu"]:
            self.gen_micro_test = self.g_builder(self.z, self.micro_coord_fake, is_training=False)
        else:
            (self.gen_micro_test, self.gen_full_test) = self.generate_full_image_gpu(self.z)

        # Patch-Guided Image Generation graph
        if self._train_content_prediction_model():
            (_, disc_real_h_rec) = self.d_builder(self.real_macro, None, is_training=False)
            estim_z = self.zp_builder(disc_real_h_rec, is_training=False)
            # I didn't especially handle this.
            # if self.config["log_params"]["merge_micro_patches_in_cpu"]:
            (_, self.rec_full) = self.generate_full_image_gpu(self.z)


        # Building these are time consuming
        if not test_mode:
            print(" [Build] Composing Loss Functions ")
            self._compose_losses()

            print(" [Build] Creating Optimizers ")
            self._create_optimizers()


    def _calc_gradient_penalty(self):
        """ Gradient Penalty for patches D """
        # This is borrowed from https://github.com/kodalinaveen3/DRAGAN/blob/master/DRAGAN.ipynb
        alpha = tf.random_uniform(shape=tf.shape(self.real_macro), minval=0.,maxval=1.)
        differences = self.gen_macro - self.real_macro # This is different from MAGAN
        interpolates = self.real_macro + (alpha * differences)
        disc_inter, _ = self.d_builder(interpolates, None, is_training=True)
        gradients = tf.gradients(disc_inter, [interpolates])[0]
        slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
        gradient_penalty = tf.reduce_mean((slopes - 1.) ** 2)
        return gradient_penalty, slopes


    def _compose_losses(self):

        # Content consistency loss
        self.code_loss = tf.reduce_mean(self.code_loss_w * tf.losses.absolute_difference(self.z, self.z_fake_pred))

        # Spatial consistency loss (reduce later)
        self.coord_mse_real = self.coord_loss_w * tf.losses.mean_squared_error(self.macro_coord_real, self.c_real_pred, reduction=NO_REDUCTION)
        self.coord_mse_fake = self.coord_loss_w * tf.losses.mean_squared_error(self.macro_coord_fake, self.c_fake_pred, reduction=NO_REDUCTION)

        # (For extrapolation training) Mask-out out-of-bound (OOB) coordinate loss since the gradients are useless
        if self.train_extrap:
            upper_bound = tf.ones([self.batch_size, self.spatial_dim], tf.float32) + 1e-4
            lower_bound = - upper_bound
            exceed_upper_bound = tf.greater(self.macro_coord_fake, upper_bound)
            exceed_lower_bound = tf.less(self.macro_coord_fake, lower_bound)

            oob_mask_sep   = tf.math.logical_or(exceed_upper_bound, exceed_lower_bound)
            oob_mask_merge = tf.math.logical_or(oob_mask_sep[:, 0], oob_mask_sep[:, 1])
            for i in range(2, self.spatial_dim):
                oob_mask_merge = tf.math.logical_or(oob_mask_merge, oob_mask_sep[:, i])
            oob_mask = tf.tile(tf.expand_dims(oob_mask_merge, 1), [1, self.spatial_dim])
            self.coord_mse_fake = tf.where(oob_mask, tf.stop_gradient(self.coord_mse_fake), self.coord_mse_fake)

        self.coord_mse_real = tf.reduce_mean(self.coord_mse_real)
        self.coord_mse_fake = tf.reduce_mean(self.coord_mse_fake)
        self.coord_loss = self.coord_mse_real + self.coord_mse_fake

        # WGAN loss
        self.adv_real = - tf.reduce_mean(self.disc_real)
        self.adv_fake = tf.reduce_mean(self.disc_fake)
        self.d_adv_loss = self.adv_real + self.adv_fake
        self.g_adv_loss = - self.adv_fake

        # Gradient penalty loss of WGAN-GP
        gradient_penalty, self.gp_slopes = self._calc_gradient_penalty()
        self.gp_loss = self.config["loss_params"]["gp_lambda"] * gradient_penalty

        # Total loss
        self.d_loss = self.d_adv_loss + self.gp_loss + self.coord_loss + self.code_loss
        self.g_loss = self.g_adv_loss + self.coord_loss + self.code_loss
        self.q_loss = self.g_adv_loss + self.code_loss

        # Wasserstein distance for visualization
        self.w_dist = - self.adv_real - self.adv_fake

        
    def _create_optimizers(self):

        t_vars = tf.trainable_variables()
        d_vars = [var for var in t_vars if 'D' in var.name]
        g_vars = [var for var in t_vars if 'G' in var.name]
        q_vars = [var for var in t_vars if 'Q' in var.name]
        
        # optimizers
        G_update_ops = tf.get_collection(self.g_builder.update_collection)
        D_update_ops = tf.get_collection(self.d_builder.update_collection)
        Q_update_ops = tf.get_collection(self.zp_builder.update_collection)
        GD_update_ops = tf.get_collection(self.cp_builder.update_collection)

        with tf.control_dependencies(G_update_ops + GD_update_ops):
            self.g_optim = tf.train.AdamOptimizer(
                self.config["train_params"]["glr"], 
                beta1=self.config["train_params"]["beta1"], 
                beta2=self.config["train_params"]["beta2"], 
            ).minimize(self.g_loss, var_list=g_vars)

        with tf.control_dependencies(D_update_ops + GD_update_ops):
            self.d_optim = tf.train.AdamOptimizer(
                self.config["train_params"]["dlr"],
                beta1=self.config["train_params"]["beta1"], 
                beta2=self.config["train_params"]["beta2"], 
            ).minimize(self.d_loss, var_list=d_vars)

        if self._train_content_prediction_model():
            with tf.control_dependencies(Q_update_ops):
                self.q_optim = tf.train.AdamOptimizer(
                    self.config["train_params"]["qlr"],
                    beta1=self.config["train_params"]["beta1"], 
                    beta2=self.config["train_params"]["beta2"], 
                ).minimize(self.q_loss, var_list=q_vars)

        if self.train_extrap:
            with tf.variable_scope("extrap_optim"):
                g_vars_partial = [
                    var for var in g_vars if ("g_resblock_0" in var.name or "g_resblock_1" in var.name)] 
                with tf.control_dependencies(G_update_ops + GD_update_ops):
                    self.g_optim_extrap = tf.train.AdamOptimizer(
                        self.config["train_params"]["glr"], 
                        beta1=self.config["train_params"]["beta1"], 
                        beta2=self.config["train_params"]["beta2"], 
                    ).minimize(self.g_loss, var_list=g_vars_partial)
    
                with tf.control_dependencies(D_update_ops + GD_update_ops):
                    self.d_optim_extrap = tf.train.AdamOptimizer(
                        self.config["train_params"]["dlr"], 
                        beta1=self.config["train_params"]["beta1"], 
                        beta2=self.config["train_params"]["beta2"], 
                    ).minimize(self.d_loss, var_list=d_vars)


    def rand_sample_full_test(self):
        if self.config["log_params"]["merge_micro_patches_in_cpu"]:
            z = self.sample_prior()
            _, full_images = self.generate_full_image_cpu(z)
        else:
            full_images = self.sess.run(
                self.gen_full_test, feed_dict={self.z: self.sample_prior()})
        return full_images

    
    def generate_full_image_gpu(self, z):
        all_micro_patches = []
        all_micro_coord = []
        num_patches_x = self.ratio_full_to_micro[0] + self.num_extrap_steps*2
        num_patches_y = self.ratio_full_to_micro[1] + self.num_extrap_steps*2
        for yy in range(num_patches_y):
            for xx in range(num_patches_x):
                if self.coordinate_system == "euclidean":
                    micro_coord_single = tf.constant([
                        self.coord_handler.euclidean_coord_int_full_to_float_micro(xx, num_patches_x, extrap_steps=self.num_extrap_steps), 
                        self.coord_handler.euclidean_coord_int_full_to_float_micro(yy, num_patches_y, extrap_steps=self.num_extrap_steps),
                    ])
                elif self.coordinate_system == "cylindrical":
                    theta_ratio = self.coord_handler.hyperbolic_coord_int_full_to_float_micro(yy, num_patches_y)
                    micro_coord_single = tf.constant([
                        self.coord_handler.euclidean_coord_int_full_to_float_micro(xx, num_patches_x), 
                        self.coord_handler.hyperbolic_theta_to_euclidean(theta_ratio, proj_func=cos),
                        self.coord_handler.hyperbolic_theta_to_euclidean(theta_ratio, proj_func=sin),
                    ])
                micro_coord = tf.tile(tf.expand_dims(micro_coord_single, 0), [tf.shape(z)[0], 1])
                generated_patch = self.g_builder(z, micro_coord, is_training=False)
                all_micro_patches.append(generated_patch)
                all_micro_coord.append(micro_coord)

        num_patches = num_patches_x * num_patches_y
        all_micro_patches = tf.concat(all_micro_patches, 0)
        all_micro_patches_reord = self.patch_handler.reord_patches_gpu(all_micro_patches, self.batch_size, num_patches)
        full_image = self.patch_handler.concat_micro_patches_gpu(
            all_micro_patches_reord, 
            ratio_over_micro=[num_patches_x, num_patches_y])

        return all_micro_patches, full_image


    def generate_full_image_cpu(self, z):
        all_micro_patches = []
        all_micro_coord = []
        num_patches_x = self.ratio_full_to_micro[0] + self.num_extrap_steps * 2
        num_patches_y = self.ratio_full_to_micro[1] + self.num_extrap_steps * 2
        for yy in range(num_patches_y):
            for xx in range(num_patches_x):
                if self.coordinate_system == "euclidean":
                    micro_coord_single = np.array([
                        self.coord_handler.euclidean_coord_int_full_to_float_micro(xx, num_patches_x, extrap_steps=self.num_extrap_steps),
                        self.coord_handler.euclidean_coord_int_full_to_float_micro(yy, num_patches_y, extrap_steps=self.num_extrap_steps),
                    ])
                elif self.coordinate_system == "cylindrical":
                    theta_ratio = self.coord_handler.hyperbolic_coord_int_full_to_float_micro(yy, num_patches_y)
                    micro_coord_single = np.array([
                        self.coord_handler.euclidean_coord_int_full_to_float_micro(xx, num_patches_x),
                        self.coord_handler.hyperbolic_theta_to_euclidean(theta_ratio, proj_func=cos),
                        self.coord_handler.hyperbolic_theta_to_euclidean(theta_ratio, proj_func=sin),
                    ])
                micro_coord = np.tile(np.expand_dims(micro_coord_single, 0), [z.shape[0], 1])
                generated_patch = self.sess.run(
                    self.gen_micro_test, feed_dict={self.z: z, self.micro_coord_fake: micro_coord}) # TODO
                all_micro_patches.append(generated_patch)
                all_micro_coord.append(micro_coord)

        num_patches = num_patches_x * num_patches_y
        all_micro_patches = np.concatenate(all_micro_patches, 0)
        all_micro_patches_reord = self.patch_handler.reord_patches_cpu(all_micro_patches, self.batch_size, num_patches)
        full_image = self.patch_handler.concat_micro_patches_cpu(
            all_micro_patches_reord, 
            ratio_over_micro=[num_patches_x, num_patches_y])

        return all_micro_patches, full_image

    def test(self, n_samples, output_dir):
        n_digits = ceil(log(n_samples, 10))
        n_iters = n_samples // self.batch_size + 1
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        for i in trange(n_iters):
            images = self.rand_sample_full_test()
            for j in range(images.shape[0]):
                global_id = i*self.batch_size + j
                if global_id < n_samples:
                    output_path = os.path.join(output_dir, "test_sample_{}.png".format(str(global_id).zfill(n_digits)))
                    imsave(output_path, images[j])
        


    def train(self, logger, evaluator, global_step):
        start_time = time.time()
        g_loss, d_loss, q_loss = 0, 0, 0
        z_fixed = self.sample_prior()
        cur_epoch = int(global_step / self.num_batches)
        cur_iter  = global_step - cur_epoch * self.num_batches

        while cur_epoch < self.epochs:
            while cur_iter < self.num_batches:

                # Create data
                z_iter = self.sample_prior()
                macro_coord, micro_coord, y_angle_ratio = self.coord_handler.sample_coord()
                feed_dict_iter = {
                    self.micro_coord_real: micro_coord,
                    self.macro_coord_real: macro_coord,
                    self.micro_coord_fake: micro_coord,
                    self.macro_coord_fake: macro_coord,
                    self.y_angle_ratio: y_angle_ratio,
                    self.z: z_iter,
                }
                feed_dict_fixed = {
                    self.micro_coord_real: micro_coord,
                    self.macro_coord_real: macro_coord,
                    self.micro_coord_fake: micro_coord,
                    self.macro_coord_fake: macro_coord,
                    self.y_angle_ratio: y_angle_ratio,
                    self.z: z_fixed,
                }
                
                # Optimize
                if (global_step % self.D_update_period) == 0:
                    _, d_summary_str, d_loss = self.sess.run(
                        [self.d_optim, logger.d_summaries, self.d_loss], 
                        feed_dict=feed_dict_iter)
                if (global_step % self.G_update_period) == 0:
                    _, g_summary_str, g_loss = self.sess.run(
                        [self.g_optim, logger.g_summaries, self.g_loss], 
                        feed_dict=feed_dict_iter)

                if self.train_extrap:
                    macro_coord_extrap, micro_coord_extrap, _ = \
                        self.coord_handler.sample_coord(num_extrap_steps=self.num_extrap_steps)
                    # Override logging inputs as well
                    feed_dict_fixed[self.micro_coord_fake] = micro_coord_extrap
                    feed_dict_fixed[self.macro_coord_fake] = macro_coord_extrap
                    feed_dict_iter[self.micro_coord_fake] = micro_coord_extrap
                    feed_dict_iter[self.macro_coord_fake] = macro_coord_extrap

                    if (global_step % self.D_update_period) == 0:
                        _, d_summary_str, d_loss = self.sess.run(
                             [self.d_optim_extrap, logger.d_summaries, self.d_loss], 
                             feed_dict=feed_dict_iter)
                    if (global_step % self.G_update_period) == 0:
                        _, g_summary_str, g_loss = self.sess.run(
                            [self.g_optim_extrap, logger.g_summaries, self.g_loss], 
                            feed_dict=feed_dict_iter)

                if self._train_content_prediction_model() and (global_step % self.Q_update_period) == 0:
                    _, q_loss = self.sess.run(
                        [self.q_optim, self.q_loss], 
                        feed_dict=feed_dict_iter)

                # Log
                time_elapsed = time.time() - start_time
                print("[{}] [Epoch: {}; {:4d}/{:4d}; global_step:{}] elapsed: {:.4f}, d: {:.4f}, g: {:.4f}, q: {:.4f}".format(
                    self.exp_name, cur_epoch, cur_iter, self.num_batches, global_step, time_elapsed, d_loss, g_loss, q_loss))
                logger.log_iter(self, evaluator, cur_epoch, cur_iter, global_step, g_summary_str, d_summary_str, 
                                z_iter, z_fixed, feed_dict_iter, feed_dict_fixed)

                cur_iter += 1
                global_step += 1
                
            cur_epoch += 1
            cur_iter = 0