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

updated pipeline, added uv_mask_loss

Browse Source
This commit is contained in:
luni
2022-06-12 15:50:30 +00:00
parent 94e693d7f6
commit def8a85cf9
6 changed files with 160 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+136 -70
View File
@@ -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