lot of things

This commit is contained in:
2022-08-06 12:27:24 +00:00
parent cd498e79f3
commit ee8f1cb68c
28 changed files with 4595 additions and 175 deletions
+15 -15
View File
@@ -17,8 +17,8 @@ from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_l
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')
# os.environ["CUDA_VISIBLE_DEVICES"] = "4"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='run_37_2', 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. . reworked data loader. per garment sum of errors. per_image_standardization for garment tempalte input. convnext discriminator, convnext generator. higher starting learning rate')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('train.py')
artifact.add_file('losses.py')
@@ -53,7 +53,7 @@ class Trainer(object):
self.train_res = t_params['train_res']
self.print_step = 100
self.save_step = 10000
self.image_summary_step = 1000
self.image_summary_step = 500
self.unsupervised_coeff = 1
self.gt_loss_coeff = 1
self.reached_max_steps = False
@@ -490,8 +490,8 @@ class Trainer(object):
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)
fake_images_05 = self.generator([test_z, test_conditions], truncation_psi=0.5, training=False)
fake_images_07 = self.generator([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)
@@ -531,12 +531,12 @@ def plot_meshes(writer, garment_unposed, garment, mesh_gt, template, body, mask,
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)):
# 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])
@@ -576,8 +576,8 @@ def main():
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,
'z_dim': 128,
'w_dim': 128,
'labels_dim': -1,
'n_mapping': 8,
'resolutions': train_resolutions,
@@ -644,7 +644,7 @@ def main():
'd_params': d_params,
# training params
'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
'g_opt': {'learning_rate': 0.001, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 1e10},
'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,
@@ -652,7 +652,7 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-35-{args["train_res"]}x{args["train_res"]}')
trainer = Trainer(training_parameters, name=f'stylegan2-test-37_2-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, static_dataset, signature_dynamic, strategy)
return