luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

plotting not working, losses fixed?

Browse Source
This commit is contained in:
luni
2022-06-10 15:32:25 +00:00
parent 31607cc754
commit 94e693d7f6
5 changed files with 108 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+56 -34
View File
@@ -5,20 +5,20 @@ import wandb
from cv2 import drawFrameAxes
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
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_ffhq import get_ffhq_dataset
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 collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_1')
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_2')
def initiate_models(g_params, d_params, use_custom_cuda):
@@ -120,54 +120,54 @@ class Trainer(object):
return params_copy
def d_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 d_tape:
# compute losses
fake_images, d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
mesh_pred = uv2mesh(fake_images)
# mesh_pred = uv2mesh(fake_images)
e_loss = edge_loss(mesh_pred, mesh_static, mask)
c_loss = collision_penalty(mesh_pred, mesh_body_posed)
# 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
# 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# + e_loss + c_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#, e_loss, c_loss
def d_train_step_reg(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 d_tape:
# compute losses
fake_images, d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
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']
mesh_pred = uv2mesh(fake_images)
# mesh_pred = uv2mesh(fake_images)
e_loss = edge_loss(mesh_pred, mesh_static, mask)
c_loss = collision_penalty(mesh_pred, mesh_body_posed)
# 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
# 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# + e_loss + c_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, d_gan_loss, r1_penalty, e_loss, c_loss
return d_loss, d_gan_loss, r1_penalty#, e_loss, c_loss
def g_train_step(self, dist_inputs):
real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
@@ -179,7 +179,7 @@ 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)
c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask)
# scale loss
g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
@@ -205,7 +205,7 @@ 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)
c_loss = collision_penalty(mesh_pred, mesh_body_posed, mask)
# scale loss
pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
@@ -224,10 +224,10 @@ class Trainer(object):
@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[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_d_loss, mean_e_loss, mean_c_loss
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
@tf.function
def dist_d_train_step_reg(inputs):
@@ -235,9 +235,9 @@ 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
# 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
@tf.function
def dist_g_train_step(inputs):
@@ -270,6 +270,8 @@ class Trainer(object):
# 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)
@@ -290,9 +292,9 @@ class Trainer(object):
# d train step
if (step + 1) % self.d_opt['reg_interval'] == 0:
d_loss, d_gan_loss, r1_penalty, e_loss, c_loss = dist_d_train_step_reg((batch, ))
d_loss, d_gan_loss, r1_penalty = dist_d_train_step_reg((batch, ))
else:
d_loss, e_loss, c_loss = dist_d_train_step((batch, ))
d_loss = dist_d_train_step((batch, ))
d_gan_loss = d_loss
r1_penalty = zero
@@ -348,6 +350,9 @@ class Trainer(object):
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)
# train_summary_writer.flush()
# print every self.print_steps
@@ -397,6 +402,23 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
return filtered_resolutions, filtered_featuremaps
def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0):
F = np.load("custom_utils/smpl_faces.npy")
garment, faces = maskFacesBatch(garment.numpy(), F, mask.astype(bool))
template, _ = maskFacesBatch(template.numpy(), F, mask.astype(bool))
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))]
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)
def main():
# global program arguments parser
parser = argparse.ArgumentParser(description='')
@@ -482,7 +504,7 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-4-{args["train_res"]}x{args["train_res"]}')
trainer = Trainer(training_parameters, name=f'stylegan2-test-5-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, strategy)
return