added skinning and gnn discriminator

This commit is contained in:
2022-06-25 16:00:58 +00:00
parent 0e88eb341c
commit 5e00b395ad
25 changed files with 321 additions and 45 deletions
+5 -3
View File
@@ -18,7 +18,7 @@ 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_28', 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')
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')
@@ -579,8 +579,10 @@ def main():
'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,
@@ -592,7 +594,7 @@ 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()
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
@@ -625,7 +627,7 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-28-{args["train_res"]}x{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