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

run-35

Browse Source
This commit is contained in:
luni
2022-06-27 12:27:55 +00:00
parent 5e00b395ad
commit cd498e79f3
11 changed files with 301 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+87 -62
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, maskFacesBatch, plot_3d
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, nest_conditions, plot_3d, batch_to_dict
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 import DatasetSimple
from dataset import DatasetSimple, DatasetFactory, StaticDataset
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
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_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.')
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')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('train.py')
artifact.add_file('losses.py')
@@ -40,7 +40,7 @@ smpl = SMPL()
class Trainer(object):
def __init__(self, t_params, name):
self.cur_tf_ver = t_params['cur_tf_ver']
# self.cur_tf_ver = t_params['cur_tf_ver']
self.use_tf_function = t_params['use_tf_function']
self.use_custom_cuda = t_params['use_custom_cuda']
self.model_base_dir = t_params['model_base_dir']
@@ -129,12 +129,13 @@ class Trainer(object):
return params_copy
def d_train_step(self, dist_inputs):
uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
real_images = uv_unposed
# uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
batch = dist_inputs[0]
real_images = batch['uv_unposed']
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)
d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions'])
# scale loss
d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
@@ -146,13 +147,14 @@ class Trainer(object):
return d_loss
def d_train_step_reg(self, dist_inputs):
uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
real_images = uv_unposed
# uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
batch = dist_inputs[0]
real_images = batch['uv_unposed']
with tf.GradientTape() as d_tape:
# compute losses
d_gan_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
self.g_params['z_dim'], conditions=conditions)
self.g_params['z_dim'], conditions=batch['conditions'])
r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']
# scale losses
@@ -168,12 +170,13 @@ class Trainer(object):
def g_train_step(self, dist_inputs):
uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
real_images = uv_unposed
# uv_posed, conditions, _, _, _, _, _, uv_unposed = dist_inputs[0]
batch = dist_inputs[0]
real_images = batch['uv_unposed']
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)
_, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions'])
# scale loss
g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
@@ -183,14 +186,15 @@ class Trainer(object):
return g_loss
def g_train_step_reg(self, dist_inputs):
uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0]
real_images = uv_unposed
# uv_posed, conditions, _, _, _, _, _, uv_unposed= dist_inputs[0]
batch = dist_inputs[0]
real_images = batch['uv_unposed']
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)
self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=batch['conditions'])
pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
# scale loss
@@ -206,28 +210,28 @@ class Trainer(object):
def g_train_step_extra(self, dist_inputs):
uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
real_images = uv_unposed
uv_mask = conditions[3]
# uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
batch = dist_inputs[0]
real_images = batch['uv_unposed']
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)
fake_images, g_loss = g_logistic_non_saturating(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=batch['conditions'])
mesh_pred_unposed = uv2mesh(fake_images)
mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)
mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed)
e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
l1_loss = 1e2 * tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask'])
c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2])
euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)
# 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
# l1_loss = tf.reduce_sum(l1_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
l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss
@@ -238,32 +242,32 @@ class Trainer(object):
def g_train_step_reg_extra(self, dist_inputs):
# real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, uv_mask = dist_inputs[0]
uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
real_images = uv_unposed
uv_mask = conditions[3]
# uv_posed, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed = dist_inputs[0]
batch = dist_inputs[0]
real_images = batch['uv_unposed']
with tf.GradientTape() as g_tape:
# compute losses
fake_images, 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)
self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=batch['conditions'])
pl_penalty = pl_penalty * self.pl_weight * self.g_opt['reg_interval']
mesh_pred_unposed = uv2mesh(fake_images)
mesh_pred = smpl.run_garment(conditions[0], conditions[1], mesh_pred_unposed)
mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed)
e_loss = 1e4 * edge_loss(mesh_pred_unposed, mesh_static, mesh_mask)
c_loss = 1e2 * collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
l1_loss = 1e2 * tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) / tf.reduce_sum(uv_mask)
euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mesh_mask[:,:,tf.newaxis] * (mesh_posed - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(mesh_mask)
e_loss = edge_loss(mesh_pred_unposed, batch['mesh_static'], batch['mesh_mask'])
c_loss = collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * batch['conditions']['uv_mask'], axis=[1,2])
euclid_error = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
ms_ssim = tf.image.ssim_multiscale(tf.transpose(real_images, [0,2,3,1]) + 1, tf.transpose(fake_images, [0,2,3,1]) + 1, max_val=2)
# 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
# l1_loss = tf.reduce_sum(l1_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
l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
# combine
@@ -273,7 +277,7 @@ class Trainer(object):
self.g_optimizer.apply_gradients(zip(g_gradients, self.generator.trainable_variables))
return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss, euclid_error, ms_ssim
def train(self, dist_datasets, strategy):
def train(self, dist_datasets: tf.data.Dataset, static_dataset: StaticDataset, dataset_signature: dict, strategy):
@tf.function
def dist_d_train_step(inputs):
per_replica_losses = strategy.run(fn=self.d_train_step, args=(inputs,))
@@ -306,9 +310,9 @@ class Trainer(object):
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.MEAN, per_replica_losses[1], axis=None)
mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[2], axis=None)
mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], 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_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim
@@ -319,9 +323,9 @@ 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)
mean_e_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[3], axis=None)
mean_c_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[4], axis=None)
mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[5], 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)
mean_l1_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[5], axis=None)
mean_euclid_error = strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses[6], axis=None)
mean_ms_ssim = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[7], axis=None)
return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss, mean_euclid_error, mean_ms_ssim
@@ -361,6 +365,7 @@ class Trainer(object):
t_start = time.time()
step = self.g_optimizer.iterations.numpy()
for batch in dist_datasets:
batch = nest_conditions(static_dataset.expand_batch(batch_to_dict(batch, dataset_signature)))
# d train step
if (step + 1) % self.d_opt['reg_interval'] == 0:
@@ -420,9 +425,6 @@ class Trainer(object):
metric_euclid_error(euclid_error)
metric_ms_ssim(ms_ssim)
# get current step
step += 1
# save to tensorboard
with train_summary_writer.as_default():
tf.summary.scalar('d_loss', metric_d_loss.result(), step=step)
@@ -447,21 +449,18 @@ class Trainer(object):
if step % self.image_summary_step == 0:
# add summary image
test_z = tf.random.normal(shape=(self.n_samples, self.g_params['z_dim']), dtype=tf.dtypes.float32)
test_conditions = batch[1] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
test_conditions = batch['conditions'] #tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
summary_image = dist_gen_samples((test_z, test_conditions))
mesh_unposed = uv2mesh(summary_image[:, :, :self.train_res, int(self.train_res/2):])
mesh = smpl.run_garment(batch[1][0], batch[1][1], mesh_unposed)
mesh = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_unposed)
# 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)
template = uv2mesh(batch[1][2])
# gt = uv2mesh(batch[0])
plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch[6], 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)
plot_meshes(mesh_writer, garment_unposed=mesh_unposed, garment=mesh, mesh_gt=batch['mesh_unposed'], template=batch['mesh_static'], body=batch['mesh_body_posed'], mask=batch['mesh_mask'].numpy(), step=step)
# train_summary_writer.flush()
@@ -479,6 +478,9 @@ class Trainer(object):
if step >= self.max_steps:
break
# get current step
step += 1
# save last checkpoint
step = self.g_optimizer.iterations.numpy()
self.manager.save(checkpoint_number=step)
@@ -492,7 +494,7 @@ class Trainer(object):
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]
template = tf.concat([test_conditions[2], test_conditions[2]], axis=2)
template = tf.concat([test_conditions['uv_static'], test_conditions['uv_static']], axis=2)
fake_image = tf.concat([fake_images_05, fake_images_07], axis=2)
final_image = tf.concat([template, fake_image], axis=3)
@@ -552,10 +554,11 @@ def main():
parser.add_argument('--train_res', default=512, type=int)
parser.add_argument('--shuffle_buffer_size', default=1000, type=int)
parser.add_argument('--batch_size_per_replica', default=4, type=int)
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
# check tensorflow version
cur_tf_ver = check_tf_version()
# cur_tf_ver = check_tf_version()
# GPU environment settings
if args['allow_memory_growth']:
@@ -594,7 +597,29 @@ 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, drop_remainder=True).repeat()
# dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').cache().batch(global_batch_size, drop_remainder=True).repeat()
factory = DatasetFactory()
signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,),
# 'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,) }
signature_static = { 'uv_static': (3, 512, 256,),
'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
'uv_mask': (512, 256,),
'uv_shaped_body': (3, 512, 256,), }
dataset = factory.get_dataset(folder='/home/workspace/data/dataset/stylegan_new/train',
batch_size=4,
signature=signature_dynamic)
static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train',
signature=signature_static)
options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
@@ -607,7 +632,7 @@ def main():
# training parameters
training_parameters = {
# global params
'cur_tf_ver': cur_tf_ver,
# 'cur_tf_ver': cur_tf_ver,
'use_tf_function': args['use_tf_function'],
'use_custom_cuda': args['use_custom_cuda'],
'model_base_dir': args['model_base_dir'],
@@ -627,8 +652,8 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-33-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, strategy)
trainer = Trainer(training_parameters, name=f'stylegan2-test-35-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, static_dataset, signature_dynamic, strategy)
return