cvpr/train.py

import os
import time
import argparse
import wandb
from cv2 import drawFrameAxes
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, plot_3d
from tensorboardX import SummaryWriter

from utils import str_to_bool
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from load_models import load_generator, load_discriminator
from dataset import DatasetSimple
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty


os.environ["CUDA_VISIBLE_DEVICES"] = "4"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_23', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, add a fading out l1 loss with gt, add tanh activation for last layer, higher weights for the extra losses, same learning rate again. added extra metrics, trained from test-22-10000')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('train.py')
artifact.add_file('losses.py')
artifact.add_file('stylegan2/layers/mapping_block.py')
run.log_artifact(artifact)

def initiate_models(g_params, d_params, use_custom_cuda):
    discriminator = load_discriminator(d_params, ckpt_dir=None, custom_cuda=use_custom_cuda)
    generator = load_generator(g_params=g_params, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
    g_clone = load_generator(g_params=g_params, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)

    # set initial g_clone weights same as generator
    g_clone.set_weights(generator.get_weights())
    return discriminator, generator, g_clone


uv2mesh = UV2Mesh(res=256)

class Trainer(object):
    def __init__(self, t_params, name):
        self.cur_tf_ver = t_params['cur_tf_ver']
        self.use_tf_function = t_params['use_tf_function']
        self.use_custom_cuda = t_params['use_custom_cuda']
        self.model_base_dir = t_params['model_base_dir']
        self.global_batch_size = t_params['batch_size']
        self.n_total_image = t_params['n_total_image']
        self.enable_unsupervised_losses = t_params['enable_unsupervised_losses']
        self.disable_l1_loss = t_params['disable_l1_loss']
        self.max_steps = int(np.ceil(self.n_total_image / self.global_batch_size))
        self.n_samples = min(t_params['batch_size'], t_params['n_samples'])
        self.train_res = t_params['train_res']
        self.print_step = 10
        self.save_step = 5000
        self.image_summary_step = 500
        self.unsupervised_coeff = 1
        self.gt_loss_coeff = 1
        self.reached_max_steps = False
        self.log_template = '{:s}, {:s}, {:s}'.format(
            'step {}: elapsed: {:.2f}s, d_loss: {:.3f}, g_loss: {:.3f}',
            'd_gan_loss: {:.3f}, g_gan_loss: {:.3f}',
            'r1_penalty: {:.3f}, pl_penalty: {:.3f}')

        # copy network params
        self.g_params = t_params['g_params']
        self.d_params = t_params['d_params']

        # set optimizer params
        self.global_batch_scaler = 1.0 / float(self.global_batch_size)
        self.r1_gamma = 10.0
        self.g_opt = self.update_optimizer_params(t_params['g_opt'])
        self.d_opt = self.update_optimizer_params(t_params['d_opt'])
        self.pl_minibatch_shrink = 2
        self.pl_weight = float(self.pl_minibatch_shrink)
        self.pl_denorm = tf.math.rsqrt(float(self.train_res) * float(self.train_res))
        self.pl_decay = 0.01
        self.pl_mean = tf.Variable(initial_value=0.0, name='pl_mean', trainable=False,
                                   synchronization=tf.VariableSynchronization.ON_READ,
                                   aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)

        # create model: model and optimizer must be created under `strategy.scope`
        self.discriminator, self.generator, self.g_clone = initiate_models(self.g_params,
                                                                           self.d_params,
                                                                           self.use_custom_cuda)

        # set optimizers
        self.d_optimizer = tf.keras.optimizers.Adam(self.d_opt['learning_rate'],
                                                    beta_1=self.d_opt['beta1'],
                                                    beta_2=self.d_opt['beta2'],
                                                    epsilon=self.d_opt['epsilon'])
        # self.d_optimizer = tf.keras.optimizers.SGD()
        self.g_optimizer = tf.keras.optimizers.Adam(self.g_opt['learning_rate'],
                                                    beta_1=self.g_opt['beta1'],
                                                    beta_2=self.g_opt['beta2'],
                                                    epsilon=self.g_opt['epsilon'])

        # setup saving locations (object based savings)
        self.ckpt_dir = os.path.join(self.model_base_dir, name)
        self.ckpt = tf.train.Checkpoint(d_optimizer=self.d_optimizer,
                                        g_optimizer=self.g_optimizer,
                                        discriminator=self.discriminator,
                                        generator=self.generator,
                                        g_clone=self.g_clone,
                                        pl_mean=self.pl_mean)
        self.manager = tf.train.CheckpointManager(self.ckpt, self.ckpt_dir, max_to_keep=20)

        # try to restore
        self.ckpt.restore(self.manager.latest_checkpoint)
        if self.manager.latest_checkpoint:
            print('Restored from {}'.format(self.manager.latest_checkpoint))

            # check if already trained in this resolution
            restored_step = self.g_optimizer.iterations.numpy()
            if restored_step >= self.max_steps:
                print('Already reached max steps {}/{}'.format(restored_step, self.max_steps))
                self.reached_max_steps = True
                return
        else:
            print('Not restoring from saved checkpoint')

    @staticmethod
    def update_optimizer_params(params):
        params_copy = params.copy()
        mb_ratio = params_copy['reg_interval'] / (params_copy['reg_interval'] + 1)
        params_copy['learning_rate'] = params_copy['learning_rate'] * mb_ratio
        params_copy['beta1'] = params_copy['beta1'] ** mb_ratio
        params_copy['beta2'] = params_copy['beta2'] ** mb_ratio
        return params_copy

    def d_train_step(self, dist_inputs):
        real_images, conditions, _, _, _, _ = dist_inputs[0]

        with tf.GradientTape() as d_tape:
            # compute losses
            d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)

            # scale loss
            d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler

            d_loss = d_loss

        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
        return d_loss

    def d_train_step_reg(self, dist_inputs):
        real_images, conditions, _, _, _, _ = dist_inputs[0]

        with tf.GradientTape() as d_tape:
            # compute losses
            d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
                                                       self.g_params['z_dim'], conditions=conditions)
            r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']

            # scale losses
            r1_penalty = tf.reduce_sum(r1_penalty) * self.global_batch_scaler
            d_gan_loss = tf.reduce_sum(d_gan_loss) * self.global_batch_scaler

            # combine
            d_loss = d_gan_loss + r1_penalty

        d_gradients = d_tape.gradient(d_loss, self.discriminator.trainable_variables)
        self.d_optimizer.apply_gradients(zip(d_gradients, self.discriminator.trainable_variables))
        return d_loss, d_gan_loss, r1_penalty


    def g_train_step(self, dist_inputs):
        real_images, conditions, _, _, _, _ = dist_inputs[0]

        with tf.GradientTape() as g_tape:
            # compute losses
            _, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)

            # scale loss
            g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
            
        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
        return g_loss

    def g_train_step_reg(self, dist_inputs):
        real_images, conditions, _, _, _, _ = dist_inputs[0]

        with tf.GradientTape() as g_tape:
            # compute losses
            _, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
                                                           self.g_params['z_dim'], self.pl_mean,
                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
            pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']

            # scale loss
            pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
            g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler
            
            # combine
            g_loss = g_gan_loss + pl_penalty

        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
        return g_loss, g_gan_loss, pl_penalty


    def g_train_step_extra(self, dist_inputs):
        real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed = dist_inputs[0]
        uv_mask = conditions[3]

        with tf.GradientTape() as g_tape:
            # compute losses
            fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)

            mesh_pred = uv2mesh(fake_images)
            e_loss = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask)
            c_loss = 1e5 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
            l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
            euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
            ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)

            # scale loss
            g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
            l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
            ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler

            g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss
            
        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
        return g_loss, e_loss, c_loss, l1_loss, euclid_error, ms_ssim

    def g_train_step_reg_extra(self, dist_inputs):
        # real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0]
        real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed = dist_inputs[0]
        uv_mask = conditions[3]

        with tf.GradientTape() as g_tape:
            # compute losses
            fake_images, g_gan_loss, pl_penalty = g_logistic_ns_pathreg(real_images, self.generator, self.discriminator,
                                                           self.g_params['z_dim'], self.pl_mean,
                                                           self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
            pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']

            mesh_pred = uv2mesh(fake_images)

            e_loss = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask)
            c_loss = 1e5 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
            l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
            euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
            ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)

            # scale loss
            pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
            g_gan_loss = tf.reduce_sum(g_gan_loss) * self.global_batch_scaler
            e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler
            c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler
            l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
            ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
            
            # combine
            g_loss = g_gan_loss + pl_penalty + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss

        g_gradients = g_tape.gradient(g_loss, self.generator.trainable_variables)
        self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
        return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss, euclid_error, ms_ssim

    def train(self, dist_datasets, strategy):
        @tf.function
        def dist_d_train_step(inputs):
            per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
            return mean_d_loss

        @tf.function
        def dist_d_train_step_reg(inputs):
            per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,))
            mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
            mean_d_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
            mean_r1_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
            return mean_d_loss, mean_d_gan_loss, mean_r1_penalty

        @tf.function
        def dist_g_train_step(inputs):
            per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
            return mean_g_loss

        @tf.function
        def dist_g_train_step_reg(inputs):
            per_replica_losses = strategy.run(fn=self.g_train_step_reg, args=(inputs,))
            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
            mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
            mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty

        @tf.function
        def dist_g_train_step_extra(inputs):
            per_replica_losses = strategy.run(fn=self.g_train_step_extra, args=(inputs,))
            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
            mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
            mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
            return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim

        @tf.function
        def dist_g_train_step_reg_extra(inputs):
            per_replica_losses = strategy.run(fn=self.g_train_step_reg_extra, args=(inputs,))
            mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[0], axis=None)
            mean_g_gan_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
            mean_pl_penalty = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[2], axis=None)
            mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[4], axis=None)
            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
            mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[6], axis=None)
            mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[7], axis=None)
            return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim

        @tf.function
        def dist_gen_samples(dist_inputs):
            per_replica_samples = strategy.run(self.gen_samples, args=(dist_inputs,))
            return per_replica_samples

        if self.reached_max_steps:
            return

        # start actual training
        print('Start Training')

        # setup tensorboards
        train_summary_writer = tf.summary.create_file_writer(self.ckpt_dir)
        mesh_writer = SummaryWriter(os.path.join(self.ckpt_dir))


        # loss metrics
        metric_d_loss = tf.keras.metrics.Mean('d_loss', dtype=tf.float32)
        metric_g_loss = tf.keras.metrics.Mean('g_loss', dtype=tf.float32)
        metric_d_gan_loss = tf.keras.metrics.Mean('d_gan_loss', dtype=tf.float32)
        metric_g_gan_loss = tf.keras.metrics.Mean('g_gan_loss', dtype=tf.float32)
        metric_r1_penalty = tf.keras.metrics.Mean('r1_penalty', dtype=tf.float32)
        metric_pl_penalty = tf.keras.metrics.Mean('pl_penalty', dtype=tf.float32)
        metric_e_loss = tf.keras.metrics.Mean('edge_loss', dtype=tf.float32)
        metric_c_loss = tf.keras.metrics.Mean('collision_loss', dtype=tf.float32)
        metric_l1_loss = tf.keras.metrics.Mean('l1_loss', dtype=tf.float32)
        metric_euclid_error = tf.keras.metrics.Mean('euclid_error', dtype=tf.float32)
        metric_ms_ssim = tf.keras.metrics.Mean('ms_ssim', dtype=tf.float32)

        # start training
        zero = tf.constant(0.0, dtype=tf.float32)
        print('max_steps: {}'.format(self.max_steps))
        t_start = time.time()
        step = self.g_optimizer.iterations.numpy()
        for batch in dist_datasets:

            # d train step
            if (step + 1) % self.d_opt['reg_interval'] == 0:
                d_loss, d_gan_loss, r1_penalty = dist_d_train_step_reg((batch, ))
            else:
                d_loss = dist_d_train_step((batch, ))
                d_gan_loss = d_loss
                r1_penalty = zero

            # g train step
            if (step + 1) % self.g_opt['reg_interval'] == 0:
                if step > self.enable_unsupervised_losses:
                    g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss, euclid_error, ms_ssim = dist_g_train_step_reg_extra((batch, ))
                else:
                    g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((batch, ))
                    e_loss = 0
                    c_loss = 0
                    l1_loss = 0
                    euclid_error = 0
                    ms_ssim = 0
            else:
                if step > self.enable_unsupervised_losses:
                    g_loss, e_loss, c_loss, l1_loss, euclid_error, ms_ssim = dist_g_train_step_extra((batch,))
                else:
                    g_loss = dist_g_train_step((batch,))
                    e_loss = 0
                    c_loss = 0
                    l1_loss = 0
                    euclid_error = 0
                    ms_ssim = 0

                g_gan_loss = g_loss
                pl_penalty = zero
                

            # if step > self.enable_unsupervised_losses:
            #     self.unsupervised_coeff = 1

            if step < self.disable_l1_loss:
                self.gt_loss_coeff = 1 - step / self.disable_l1_loss
            else:
                self.gt_loss_coeff = 0
            
            # update g_clone
            self.g_clone.set_as_moving_average_of(self.generator)

            # update metrics
            metric_d_loss(d_loss)
            metric_g_loss(g_loss)
            metric_d_gan_loss(d_gan_loss)
            metric_g_gan_loss(g_gan_loss)
            metric_r1_penalty(r1_penalty)
            metric_pl_penalty(pl_penalty)
            metric_e_loss(e_loss)
            metric_c_loss(c_loss)
            metric_l1_loss(l1_loss)
            metric_euclid_error(euclid_error)
            metric_ms_ssim(ms_ssim)

            # get current step
            step += 1

            # save to tensorboard
            with train_summary_writer.as_default():
                tf.summary.scalar('d_loss', metric_d_loss.result(), step=step)
                tf.summary.scalar('g_loss', metric_g_loss.result(), step=step)
                tf.summary.scalar('d_gan_loss', metric_d_gan_loss.result(), step=step)
                tf.summary.scalar('g_gan_loss', metric_g_gan_loss.result(), step=step)
                tf.summary.scalar('r1_penalty', metric_r1_penalty.result(), step=step)
                tf.summary.scalar('pl_penalty', metric_pl_penalty.result(), step=step)
                tf.summary.scalar('edge_loss', metric_e_loss.result(), step=step)
                tf.summary.scalar('collision_loss', metric_c_loss.result(), step=step)
                tf.summary.scalar('unsupervised_coeff', self.unsupervised_coeff, step=step)
                tf.summary.scalar('gt_loss_coeff', self.gt_loss_coeff, step=step)
                tf.summary.scalar('l1_loss', metric_l1_loss.result(), step=step)
                tf.summary.scalar('euclid_error', metric_euclid_error.result(), step=step)
                tf.summary.scalar('ms_ssim', metric_ms_ssim.result(), step=step)

            # save every self.save_step
            if step % self.save_step == 0:
                self.manager.save(checkpoint_number=step)

            # save every self.image_summary_step
            if step % self.image_summary_step == 0:
                # add summary image
                test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32)
                test_conditions = batch[1] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
                summary_image = dist_gen_samples((test_z, test_conditions))

                mesh = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):])
                # convert to tensor image
                summary_image = self.convert_per_replica_image(summary_image, strategy)

                with train_summary_writer.as_default():
                    tf.summary.image('images', summary_image, step=step)

                template = uv2mesh(batch[1][2])
                # gt = uv2mesh(batch[0])
                plot_meshes(mesh_writer, garment=mesh, mesh_gt=batch[5], template=template, body=batch[3], mask=batch[4].numpy(), step=step)
                # plot_meshes(mesh_writer, garment=gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step)

                # train_summary_writer.flush()

            # print every self.print_steps
            if step % self.print_step == 0:
                elapsed = time.time() - t_start
                print(self.log_template.format(step, elapsed, d_loss.numpy(), g_loss.numpy(),
                                               d_gan_loss.numpy(), g_gan_loss.numpy(),
                                               r1_penalty.numpy(), pl_penalty.numpy()))

                # reset timer
                t_start = time.time()

            # check exit status
            if step >= self.max_steps:
                break

        # save last checkpoint
        step = self.g_optimizer.iterations.numpy()
        self.manager.save(checkpoint_number=step)
        return

    def gen_samples(self, inputs):
        test_z, test_conditions = inputs

        # run networks
        fake_images_05 = self.g_clone([test_z, test_conditions], truncation_psi=0.5, training=False)
        fake_images_07 = self.g_clone([test_z, test_conditions], truncation_psi=0.7, training=False)

        # merge on batch dimension: [n_samples, 3, out_res, 2 * out_res]
        template = tf.concat([test_conditions[2], test_conditions[2]], axis=2)
        fake_image = tf.concat([fake_images_05, fake_images_07], axis=2)
        final_image = tf.concat([template, fake_image], axis=3)

        return final_image

    @staticmethod
    def convert_per_replica_image(nchw_per_replica_images, strategy):
        as_tensor = tf.concat(strategy.experimental_local_results(nchw_per_replica_images), axis=0)
        as_tensor = tf.transpose(as_tensor, perm=[0, 2, 3, 1])
        as_tensor = (tf.clip_by_value(as_tensor, -1.0, 1.0) + 1.0) * 127.5
        as_tensor = tf.cast(as_tensor, tf.uint8)
        return as_tensor


def filter_resolutions_featuremaps(resolutions, featuremaps, res):
    index = resolutions.index(res)
    filtered_resolutions = resolutions[:index + 1]
    filtered_featuremaps = featuremaps[:index + 1]
    return filtered_resolutions, filtered_featuremaps


def plot_meshes(writer, garment, mesh_gt, template, body, mask, step, nr_offset=0):
    body = body.numpy()
    template = template.numpy()
    mesh_gt = mesh_gt.numpy()
    F = np.load("custom_utils/smpl_faces.npy")
    garment, faces = maskFacesBatch(garment.numpy(), F, mask.astype(bool))
    template = [template[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
    mesh_gt = [mesh_gt[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]

    collision = [np.concatenate((body[i], garment[i]), axis=0) for i in range(len(garment))]
    collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment))]
    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment[i], dtype=np.int32)*230), axis=0) for i in range(len(garment))]

    for i in range(len(garment)):
        writer.add_mesh('garment_pc_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], global_step=step)
        writer.add_mesh('garment_pred_' + str(i+nr_offset), vertices=garment[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
        writer.add_mesh('template_' + str(i+nr_offset), vertices=template[i][np.newaxis,:,:], faces=faces[i][np.newaxis,:,:], global_step=step)
        writer.add_mesh('body_' + str(i+nr_offset), vertices=body[i][np.newaxis,:,:], faces=F[np.newaxis,:,:], global_step=step)
        writer.add_mesh('collision_' + str(i+nr_offset), vertices=collision[i][np.newaxis,:,:], faces=collision_faces[i][np.newaxis,:,:], global_step=step, colors=collision_color[i][np.newaxis,:,:])
    
    for i in range(len(garment)):
        plotly_fig = plot_3d(mesh_gt[i], garment[i], faces[i], collision[i], collision_faces[i], collision_color[i])
        wandb.log({"garment_" + str(i): plotly_fig})

def main():
    # global program arguments parser
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('--allow_memory_growth', type=str_to_bool, nargs='?', const=True, default=True)
    parser.add_argument('--debug_split_gpu', type=str_to_bool, nargs='?', const=True, default=False)
    parser.add_argument('--use_tf_function', type=str_to_bool, nargs='?', const=True, default=True)
    parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=True)
    parser.add_argument('--model_base_dir', default='./models', type=str)
    parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
    parser.add_argument('--train_res', default=512, type=int)
    parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
    parser.add_argument('--batch_size_per_replica', default=4, type=int)
    args = vars(parser.parse_args())

    # check tensorflow version
    cur_tf_ver = check_tf_version()

    # GPU environment settings
    if args['allow_memory_growth']:
        allow_memory_growth()
    if args['debug_split_gpu']:
        split_gpu_for_testing(mem_in_gb=4.5)

    # network params
    resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
    g_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8]
    d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
    # featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
    # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
    train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, args['train_res'])
    _, train_d_featuremaps = filter_resolutions_featuremaps(resolutions, d_featuremaps, args['train_res'])
    g_params = {
        'z_dim': 256,
        'w_dim': 256,
        'labels_dim': -1,
        'n_mapping': 8,
        'resolutions': train_resolutions,
        'featuremaps': train_g_featuremaps,
    }
    d_params = {
        'labels_dim': -1,
        'resolutions': train_resolutions,
        'featuremaps': train_d_featuremaps,
    }

    # prepare distribute strategy
    strategy = tf.distribute.MirroredStrategy()
    global_batch_size = args['batch_size_per_replica'] * strategy.num_replicas_in_sync

    # prepare dataset
    # dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
    dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size).repeat()

    options = tf.data.Options()
    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
    dataset = dataset.with_options(options)

    with strategy.scope():
        # distribute dataset        
        dist_dataset = strategy.experimental_distribute_dataset(dataset)

        # training parameters
        training_parameters = {
            # global params
            'cur_tf_ver': cur_tf_ver,
            'use_tf_function': args['use_tf_function'],
            'use_custom_cuda': args['use_custom_cuda'],
            'model_base_dir': args['model_base_dir'],
            'enable_unsupervised_losses': 0,
            'disable_l1_loss': 10000,

            # network params
            'g_params': g_params,
            'd_params': d_params,

            # training params
            'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
            'd_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 16},
            'batch_size': global_batch_size,
            'n_total_image': 25000000,
            'n_samples': 12,
            'train_res': args['train_res'],
        }

        trainer = Trainer(training_parameters, name=f'stylegan2-test-23-{args["train_res"]}x{args["train_res"]}')
        trainer.train(dist_dataset, strategy)
    return


if __name__ == '__main__':
    main()