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
from custom_utils.smpl import SMPL

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_33', 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, lower discr. features too. use reduce_sum for losses divided by sum of mask. learn unposed version then do skinning. lowered collsion coeff. learns offsets over zero body. fixed unposed uv map conversion. gnn discriminator.')
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)
smpl = SMPL()

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 = 100
        self.save_step = 10000
        self.image_summary_step = 1000
        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):
        uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
        real_images = uv_unposed

        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):
        uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
        real_images = uv_unposed

        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):
        uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
        real_images = uv_unposed

        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):
        uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0]
        real_images = uv_unposed

        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):
        uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
        real_images = uv_unposed
        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_unposed = uv2mesh(fake_images)
            mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)

            e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
            c_loss = 1e2 * 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]
        uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
        real_images = uv_unposed
        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_unposed = uv2mesh(fake_images)
            mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)

            e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
            c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
            l1_loss = 1e2 * tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) / tf.reduce_sum(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.MEAN, per_replica_losses[1], axis=None)
            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[2], axis=None)
            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, 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.MEAN, per_replica_losses[3], axis=None)
            mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
            mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, 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_unposed = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):])
                mesh = smpl.run_garment(batch[1][0], batch[1][1], mesh_unposed)
                # 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_unposed=mesh_unposed, garment=mesh, mesh_gt=batch[6], 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_unposed, garment, mesh_gt, template, body, mask, step, nr_offset=0):
    body = body.numpy()
    template = template.numpy()
    mesh_gt = mesh_gt.numpy()
    garment_unposed = garment_unposed.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))]
    garment_unposed = [garment_unposed[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_unposed[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 = [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,
        'zero_body': np.load('custom_utils/uv_zero_body.npy')
    }
    d_params = {
        'batch_size': args['batch_size_per_replica'],
        '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, drop_remainder=True).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': 30000,

            # 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-33-{args["train_res"]}x{args["train_res"]}')
        trainer.train(dist_dataset, strategy)
    return


if __name__ == '__main__':
    main()