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, nest_conditions, plot_3d, batch_to_dict 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, DatasetFactory, StaticDataset 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_35', 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 and edge coeff. learns offsets over shaped body. fixed unposed uv map conversion. gnn discriminator. reworked data loader. per garment sum of errors. per_image_standardization for garment tempalte input') 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] batch = dist_inputs[0] real_images = batch['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=batch['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] batch = dist_inputs[0] real_images = batch['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=batch['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] batch = dist_inputs[0] real_images = batch['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=batch['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] batch = dist_inputs[0] real_images = batch['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=batch['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] batch = dist_inputs[0] real_images = batch['uv_unposed'] 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=batch['conditions']) mesh_pred_unposed = uv2mesh(fake_images) mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed) e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask']) c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']) l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2]) euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['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] batch = dist_inputs[0] real_images = batch['uv_unposed'] 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=batch['conditions']) pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval'] mesh_pred_unposed = uv2mesh(fake_images) mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed) e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask']) c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']) l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2]) euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['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: tf.data.Dataset, static_dataset: StaticDataset, dataset_signature: dict, 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: batch = nest_conditions(static_dataset.expand_batch(batch_to_dict(batch, dataset_signature))) # 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) # 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['conditions'] #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['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], 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) plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch['mesh_unposed'], template=batch['mesh_static'], body=batch['mesh_body_posed'], mask=batch['mesh_mask'].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 # get current step step += 1 # 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['uv_static'], test_conditions['uv_static']], 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) parser.add_argument('--name', type=str) 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() factory = DatasetFactory() signature_dynamic = { 'id': (), 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), # 'uv_body_posed': (3, 512, 256,), 'mesh_body_posed': (14475, 3,), 'mesh_posed': (14475, 3,), 'mesh_unposed': (14475, 3,), 'uv_unposed': (3, 512, 256,) } signature_static = { 'uv_static': (3, 512, 256,), 'mesh_static': (14475, 3,), 'mesh_mask': (14475,), 'uv_mask': (512, 256,), 'uv_shaped_body': (3, 512, 256,), } dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train', batch_size=4, signature=signature_dynamic) static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train', signature=signature_static) 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-35-{args["train_res"]}x{args["train_res"]}') trainer.train(dist_dataset, static_dataset, signature_dynamic, strategy) return if __name__ == '__main__': main()