From def8a85cf9c0eb8bb3a71fda6465a828757a5bef Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 12 Jun 2022 15:50:30 +0000 Subject: [PATCH] updated pipeline, added uv_mask_loss --- custom_utils/custom_utils.py | 3 +- custom_utils/mask_256x512.npy | Bin 0 -> 131200 bytes dataset.py | 7 +- inference_custom.py | 11 +- losses.py | 11 ++ train.py | 206 ++++++++++++++++++++++------------ 6 files changed, 160 insertions(+), 78 deletions(-) create mode 100644 custom_utils/mask_256x512.npy diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index aa151e8..37d13a7 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -21,7 +21,7 @@ class UV2Mesh(): super(UV2Mesh, self).__init__() self.res = res uv_map = np.load('custom_utils/stretched_uv_map_3.npy') - self.uv_map_discreet = np.round(uv_map * self.res).astype(np.int32) + self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32) m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item() m2uv = m2uv_array(m2uv) self.uv_pixels = self.uv_map_discreet[m2uv] @@ -35,6 +35,7 @@ class UV2Mesh(): # self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1) # input = channels_first + @tf.function def uv_to_mesh(self, uv): uv = tf.transpose(uv, [0,2,3,1]) batch_size = uv.shape[0] diff --git a/custom_utils/mask_256x512.npy b/custom_utils/mask_256x512.npy new file mode 100644 index 0000000000000000000000000000000000000000..770d714093da9b9dacced063eb3ba446e67a53ef GIT binary patch literal 131200 zcmeI4yRsZf8H6R5r#PMF?>>Bd`tX;Z zzWnpcr@#O4r!PN#exCpJr@wvue4GFJ*H8cWd>jA#hwr}s`1Jh`zx>6=r~f`x{rTBD z|KkmK1Kxl)aBtwf_RrlW_EYQpK4!U}8wWAYIc#AZ#W-irA2cxRHViWk!rv~3VQ7tQ z$b*p<Sc@Jy9Hf@LDl1j>*Hp!D^h3YwZ905iZcm0tWiP~7(7`KOAn zeE~%8zvmY=5N-T$QV0|O2Bb7ROd4NQ;zYlUC=KHgk7s-H^!7VW-~9nNTYTwPz)1Jg zn9xD(40Tph=HEeiepLgzms-JuifR8Fa6hMl_`oGO1% zal@saPVdkVzr!2A^Gh0Nx#{6q30&|Sph!P>z=UbauSI+dh)93-{0pDwmo%`vzM3y2 zcdwFFCPRk~dAS0xpvN|(;B`J57c{Ndqmf|460ljR%m{DKBFuPVXgU-bqMbb9~q3iMk* zf+`WElfUXoP$iREebL*EArxSprc_`p0<#Ts*%l{)0RKJ-@60 z!NZD|l%;=+Zvevh5Agu-{K5wKcPp+aO#cWE0BQUOcmQ~QWdr!T6&F;dziH#VwBW__ zi{np!|2@CB0s8$4>xt9P{s5H6pZo#f`Q;73Z&+4No__8JK!^BKKLC2Nwy48I7Yo(2CHC8JuiJy5H3TKPSK#3K}&CGpcNgZ6K?fH*dp02ebrdR)?KzN ze0RK_hN&ugGECz2T*V|Pj^Pb>1Kxl);0<^K-hemY4R{0IfH&X`cmv*mH{cC;1Kxl) z;0<^K-hemY4R{0IfH&X`cmv*mH{cC;1Kxl);0<^K-hemY4R{0IfH&X`cmv*mH{cC; z1Kxl);0>4tuKIW2;(ra&RPzlE;;Cb^-&FWmz}&gVEh0Bz_P)ocB0Em)#%XJ)4ODY( zt`cez5N9T6fr!zjF|33(Z1{B77GQQxPTvB@DpM*bq39o=kO3=%b4qyi&Df~`?4o&; zZ3W2NweEkkt5`mp;e3bP|9~SCanFCUyg%@eGp{M*=r3#WC7<;Mya8{(8}J6a0dK$? z@CLjAZ@?Sy2Ji;VZ{qymh?jRWibvM#nmNJ`R-@e&Z)OyCe(?!&h9AUDUuJnTqj=yP zE}1j@AhgTTVYvE7D`7N64#+}a2x?aeNHkhI+|;NC`X@Ao!icHx7hZDavAH-ZtvmnpLe zAEaDk^SHSYj5$8fnN9d0@iD?TAk4AW2 zJ%6I})HHm@H*#^;%_epjw@lBn?i;nZ?Pe1@j9aE>S@(@v+;+2x9mXxwv#k3@Eo!^j z!w#d4hm)+oRf`(iE$}GncsL0I7`dpi*+CDYj)#*pfzgXOn+@x zVt!|HEP+*vdD|WADCT!I#}Zhzh_~Iqjv{vFQ!Ij2i@4hj>?mTFF~u?%wTQb}hX)b6 zj42qx$nEXTGCbJ3<;=hkMsDwKmf6AHEN_OTFlu`Mv&;_mW_dF#g;CoBm}PdbH_Mx0 zDU8| z4ad4bUu}#EhcwNqXF=9^`pWv0j%b=y&yuY3^p*7~9nmzao+Vl5;VWxXK7w&pIg7H+ z!&la(d<5gHat39crmu`m#en8n^bE@~OJ5k5ssT;3=!G=PBz+N@)Lqj=i(ZJcOwt#j zN!>L~wCII6%OrdedemOSNQ+#^vrNJlp-1gCjI_uFJj)z?5qMNz(oA=~0O&r2{|r2; zFL9{5;7W9#p?|IpEmt(vU2kQ&&(J?thn6dv>aMpk-Dl{Zt3%5bO?B5>neH?2&()yq z3dXw2El~Ff`sZTMctLZ$=`BMjU(kKT}E(rQ-Sb_-LmTQ6qj7r+iJxsMIb`8b6_A zRBFf5l+P*|mD=S=<0piSLg{#xs#zVSPPr;*;*gF~r<{mVHLau6DOUwe9g<)^o5_Sz zN$Tchm@?tIoT*bPOo=Y~<^Q!26{ku!bsuZzRh%l_+MLOA63~Zv{K-H?wI*+4^@0dg zBwHObc~Ao?Q58J?Z?$Mp6{^X5*tDP^)o3=wOdi&-DpZpZnidqK8qKDd>BAaWUNxPd zZG|V*D7K|cpVrXws_6u6D?F)2u`T7nX$>x}dXS)P2`7~Zwq-mzEy2Y_j}o-4;iwY9 zwv30TCA_rgVS=_r8dag!mhkwrf|u4jPSCbWvnur35}uq^0K=Lm3EGxnR)JhwzO&N; zU|8}jLEAbEE0Al;cXnC;3`?FRXk3SBdAY`XXUBP9R`HC#b)BW<)mrnNo#%;J#WViq zbru&_YtDBjfD@yFXY}psC@wD6p6^TnM@9wD=z(>Vmlgx^ok`Hhq}~}ju}<>RT1386 z4H}x$J4Ht}aJ;k<67O7u#>Vu{(UBb(FD!(F%MpRODY+bUWDf<*>LAhbL|AZ0Ee{>p z1reh%NU%&16dV%EL}#`^%%BXDD^~>shoo}Jp?wfFCj;fmmEph{rCe@oAA^mlV7YRI zFmgmFmm1uMa8n{+u3RY!ozTf;CifBClm?kATZ@7Rbh4SrjQ}?!L8dCug6P|23Yf{Q z?cj_8nW;#N!EaV6VkS2?lQRNjq#`W_zg49O8Qt7U&YVXh6^SwYttv&x?B-T<b?pAi>JR7M<4dHJUDI!KUw~|BW(M%;`5Pz#o2{O62l^Z)xhRWB% z@SBD5naQ2a*qi_vD^rVur<5|mtG7{{eRQg0h6nqg>RrN2=hu3kfexwo1Cd5qKUXoz z+kL$OZ@?Sy2D|}pz#H%eya8{(8}J6a0dK$?@CLjAZ@?Sy2D|}pz#H%eya8{(8}J6a z0dK$?@CLjAZ@?Sy2D|}pz#H%eya8{(8}J6a0dK$?@CLjAZ@?Sy2E2iz2Cn*d;NpJ` z@~EkAVJ)6IHv284T{6T1=FTl54|kE9Fnix)G>{#qcH^`))CQ_KH&+QY35YWjv_QmY z(->Al8#a8pYYQ+tC#P?LW0fftlu+~!P{@E4!Z{_p`ey7@0Cv$l%C-XJ?OOLg+Epx{ w&2YZM?tj3MiMZ!KS>7Lb$eGuaarBq9_>#|h1Kxl);0<^K-hemY4NwjI9~iE6W&i*H literal 0 HcmV?d00001 diff --git a/dataset.py b/dataset.py index d3e6c9f..58ac10b 100644 --- a/dataset.py +++ b/dataset.py @@ -23,15 +23,16 @@ class DatasetSimple(tf.data.Dataset): uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]) mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static'] mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) - mask = np.load(os.path.join(folder, 'mask', id + '.npy')); mask = mask.astype(np.float32) + mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32) + uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32) # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask - yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mask + yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static), mesh_static, mesh_body_posed, mesh_mask, uv_mask def __new__(cls, folder, res=256): return tf.data.Dataset.from_generator( cls._generator, - output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32) ), + output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32) ), args=(folder, ) ) diff --git a/inference_custom.py b/inference_custom.py index 08788b2..ccfa0f0 100644 --- a/inference_custom.py +++ b/inference_custom.py @@ -9,6 +9,7 @@ from dataset import DatasetSimple from load_models import load_generator from tf_utils import allow_memory_growth +os.environ["CUDA_VISIBLE_DEVICES"] = "3" def m2uv_array(m2uv): result = np.zeros((14475, 2), dtype=np.int32) @@ -39,7 +40,7 @@ def uv_to_mesh(uv): return mesh -def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100): +def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=50): # create generator resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] @@ -53,7 +54,9 @@ def inference(ckpt_dir, use_custom_cuda, res, folder, data_folder, samples=100): 'featuremaps': featuremaps[:filter_index + 1], } generator = load_generator(g_params, is_g_clone=True, ckpt_dir=ckpt_dir, custom_cuda=use_custom_cuda) - dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(1).repeat() + dataset = DatasetSimple(data_folder).cache().batch(1).repeat() + + os.makedirs(os.path.join(folder), exist_ok=True) i = 0 for batch in dataset: @@ -80,9 +83,9 @@ def main(): checkpoints = [ { 'res': 512, - 'ckpt_dir': 'models/stylegan2-test-5-512x512', + 'ckpt_dir': 'models/stylegan2-test-8-512x512', 'use_custom_cuda': True, - 'folder': 'output', + 'folder': 'output/test-8-196000', 'data_folder': '/home/workspace/data/dataset/stylegan_new/train' }, ] diff --git a/losses.py b/losses.py index c924b7a..ee5bb63 100644 --- a/losses.py +++ b/losses.py @@ -187,6 +187,17 @@ def collision_penalty(va, vb, mask, eps=1e-3): return tf.reduce_sum(interpenetration**3) / batch_size +whole_garment_uv_mask = np.load('custom_utils/mask_256x512.npy') + +def uv_mask_loss(pred, uv_mask): + inv_mask = (1 - uv_mask) * whole_garment_uv_mask + pred = tf.reduce_sum(pred, axis=1) + invalid = pred * inv_mask + invalid = invalid > 0.01 + invalid_count = tf.math.count_nonzero(invalid, dtype=tf.float32) + total_count = tf.math.count_nonzero(inv_mask, dtype=tf.float32) + return invalid_count / total_count + if __name__ == '__main__': V = tf.ones((4, 14475, 3)) diff --git a/train.py b/train.py index 1d56fc2..76ccefb 100644 --- a/train.py +++ b/train.py @@ -12,14 +12,17 @@ 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 +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"] = "1" -wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_2') - +os.environ["CUDA_VISIBLE_DEVICES"] = "2" +run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_7(test-10)') +artifact = wandb.Artifact('code', type='code') +artifact.add_file('train.py') +artifact.add_file('losses.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) @@ -41,18 +44,19 @@ class Trainer(object): 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.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 = 1000 + self.save_step = 5000 self.image_summary_step = 500 + self.unsupervised_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}', - 'e_loss: {:.3f}, c_loss: {:.3f}') + 'r1_penalty: {:.3f}, pl_penalty: {:.3f}') # copy network params self.g_params = t_params['g_params'] @@ -94,7 +98,7 @@ class Trainer(object): 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=5) + self.manager = tf.train.CheckpointManager(self.ckpt, self.ckpt_dir, max_to_keep=20) # try to restore self.ckpt.restore(self.manager.latest_checkpoint) @@ -120,30 +124,24 @@ class Trainer(object): return params_copy def d_train_step(self, dist_inputs): - real_images, conditions, _, _, _ = dist_inputs[0] + 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) - # mesh_pred = uv2mesh(fake_images) - - # e_loss = edge_loss(mesh_pred, mesh_static, mask) - # c_loss = collision_penalty(mesh_pred, mesh_body_posed) # scale loss d_loss = tf.reduce_sum(d_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 - d_loss = d_loss# + e_loss + c_loss + 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#, e_loss, c_loss + return d_loss def d_train_step_reg(self, dist_inputs): - real_images, conditions, _, _, _ = dist_inputs[0] + real_images, conditions, _, _, _, _ = dist_inputs[0] with tf.GradientTape() as d_tape: # compute losses @@ -151,49 +149,80 @@ class Trainer(object): self.g_params['z_dim'], conditions=conditions) r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval'] - # mesh_pred = uv2mesh(fake_images) - - # e_loss = edge_loss(mesh_pred, mesh_static, mask) - # c_loss = collision_penalty(mesh_pred, mesh_body_posed) - # 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 - # e_loss = tf.reduce_sum(e_loss) * self.global_batch_scaler - # c_loss = tf.reduce_sum(c_loss) * self.global_batch_scaler # combine - d_loss = d_gan_loss + r1_penalty# + e_loss + c_loss + 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#, e_loss, c_loss + return d_loss, d_gan_loss, r1_penalty + def g_train_step(self, dist_inputs): - real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0] + 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, uv_mask = dist_inputs[0] 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 = edge_loss(mesh_pred, mesh_static, mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask) + e_loss = 0.1 * edge_loss(mesh_pred, mesh_static, mesh_mask) + c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) + m_loss = uv_mask_loss(fake_images, uv_mask) # 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 + m_loss = tf.reduce_sum(m_loss) * self.global_batch_scaler - g_loss = g_loss + e_loss + c_loss + g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss + m_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 + return g_loss, e_loss, c_loss, m_loss - def g_train_step_reg(self, dist_inputs): - real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0] + def g_train_step_reg_extra(self, dist_inputs): + real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0] with tf.GradientTape() as g_tape: # compute losses @@ -204,30 +233,30 @@ class Trainer(object): mesh_pred = uv2mesh(fake_images) - e_loss = edge_loss(mesh_pred, mesh_static, mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask) + e_loss = 0.1 * edge_loss(mesh_pred, mesh_static, mesh_mask) + c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) + m_loss = uv_mask_loss(fake_images, uv_mask) # 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 + m_loss = tf.reduce_sum(m_loss) * self.global_batch_scaler # combine - g_loss = g_gan_loss + pl_penalty + e_loss + c_loss + g_loss = g_gan_loss + pl_penalty + self.unsupervised_coeff * (e_loss + c_loss + m_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 + return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, m_loss 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) - # 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) - return mean_d_loss#, mean_e_loss, mean_c_loss + return mean_d_loss @tf.function def dist_d_train_step_reg(inputs): @@ -235,17 +264,13 @@ class Trainer(object): 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) - # 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) - return mean_d_loss, mean_d_gan_loss, mean_r1_penalty#, mean_e_loss, mean_c_loss + 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[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) - return mean_g_loss, mean_e_loss, mean_c_loss + 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): @@ -253,9 +278,27 @@ class Trainer(object): 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_m_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None) + return mean_g_loss, mean_e_loss, mean_c_loss, mean_m_loss + + @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) - return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss + mean_m_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None) + return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_m_loss @tf.function def dist_gen_samples(dist_inputs): @@ -282,6 +325,7 @@ class Trainer(object): 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_m_loss = tf.keras.metrics.Mean('uv_mask_loss', dtype=tf.float32) # start training zero = tf.constant(0.0, dtype=tf.float32) @@ -300,12 +344,29 @@ class Trainer(object): # g train step if (step + 1) % self.g_opt['reg_interval'] == 0: - g_loss, g_gan_loss, pl_penalty, e_loss, c_loss = dist_g_train_step_reg((batch, )) + if step > self.enable_unsupervised_losses: + g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, m_loss = 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 + m_loss = 0 else: - g_loss, e_loss, c_loss = dist_g_train_step((batch,)) + if step > self.enable_unsupervised_losses: + g_loss, e_loss, c_loss, m_loss = dist_g_train_step_extra((batch,)) + else: + g_loss = dist_g_train_step((batch,)) + e_loss = 0 + c_loss = 0 + m_loss = 0 + g_gan_loss = g_loss pl_penalty = zero + + # if step > self.enable_unsupervised_losses: + # self.unsupervised_coeff = (step / 40000) + # update g_clone self.g_clone.set_as_moving_average_of(self.generator) @@ -318,6 +379,7 @@ class Trainer(object): metric_pl_penalty(pl_penalty) metric_e_loss(e_loss) metric_c_loss(c_loss) + metric_m_loss(m_loss) # get current step step += 1 @@ -332,6 +394,8 @@ class Trainer(object): 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('uv_mask_loss', metric_m_loss.result(), step=step) # save every self.save_step if step % self.save_step == 0: @@ -344,14 +408,17 @@ class Trainer(object): 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) - # mesh = uv2mesh(tf.transpose(summary_image[:, :self.train_res, :, :], [0,3,1,2])) - # plot_meshes(mesh_writer, garment=mesh, template=batch[1][2], body=batch[3], mask=batch[4], step=step) + template = uv2mesh(batch[1][2]) + # gt = uv2mesh(batch[0]) + plot_meshes(mesh_writer, garment=mesh, 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() @@ -360,8 +427,7 @@ class Trainer(object): 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(), - e_loss.numpy(), c_loss.numpy())) + r1_penalty.numpy(), pl_penalty.numpy())) # reset timer t_start = time.time() @@ -376,14 +442,17 @@ class Trainer(object): return def gen_samples(self, inputs): - test_z, test_labels = inputs + test_z, test_conditions = inputs # run networks - fake_images_05 = self.g_clone([test_z, test_labels], truncation_psi=0.5, training=False) - fake_images_07 = self.g_clone([test_z, test_labels], truncation_psi=0.7, training=False) + 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] - final_image = tf.concat([fake_images_05, fake_images_07], axis=2) + 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 @@ -409,14 +478,15 @@ def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0): body = body.numpy() collision = [np.concatenate((body[i], garment[i]), axis=0) for i in range(len(garment))] - collision_faces = [np.concatenate((F, faces + 14475), 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[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,:,:]) def main(): @@ -444,7 +514,8 @@ def main(): # network params resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] - featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] + featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16] + # featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32] # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4] train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res']) g_params = { @@ -468,12 +539,6 @@ def main(): # 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() - # limit = int(25594/global_batch_size) - # count = 0 - # for d in dataset: - # count += global_batch_size - # if count > limit: - # break options = tf.data.Options() options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF @@ -490,6 +555,7 @@ def main(): '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, # network params 'g_params': g_params, @@ -504,7 +570,7 @@ def main(): 'train_res': args['train_res'], } - trainer = Trainer(training_parameters, name=f'stylegan2-test-5-{args["train_res"]}x{args["train_res"]}') + trainer = Trainer(training_parameters, name=f'stylegan2-test-10-{args["train_res"]}x{args["train_res"]}') trainer.train(dist_dataset, strategy) return