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

prototype conditional pipeline

Browse Source
This commit is contained in:
luni
2022-06-10 06:59:58 +00:00
parent 03c22c1ba6
commit 31607cc754
24 changed files with 1093 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+98 -47
View File
@@ -5,13 +5,16 @@ import wandb
from cv2 import drawFrameAxes
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh
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
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"
@@ -28,6 +31,8 @@ def initiate_models(g_params, d_params, use_custom_cuda):
return discriminator, generator, g_clone
uv2mesh = UV2Mesh(res=256)
class Trainer(object):
def __init__(self, t_params, name):
self.cur_tf_ver = t_params['cur_tf_ver']
@@ -41,12 +46,13 @@ class Trainer(object):
self.train_res = t_params['train_res']
self.print_step = 10
self.save_step = 1000
self.image_summary_step = 1000
self.image_summary_step = 500
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}')
'r1_penalty: {:.3f}, pl_penalty: {:.3f}',
'e_loss: {:.3f}, c_loss: {:.3f}')
# copy network params
self.g_params = t_params['g_params']
@@ -114,85 +120,114 @@ class Trainer(object):
return params_copy
def d_train_step(self, dist_inputs):
real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
with tf.GradientTape() as d_tape:
# compute losses
d_loss, cond_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], labels=condition)
fake_images, 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
cond_loss = tf.reduce_sum(cond_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 + cond_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, cond_loss
return d_loss, e_loss, c_loss
def d_train_step_reg(self, dist_inputs):
real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
real_images, conditions, mesh_static, mesh_body_posed, mask = dist_inputs[0]
with tf.GradientTape() as d_tape:
# compute losses
d_gan_loss, cond_loss, r1_penalty = d_logistic_r1_reg(real_images, self.generator, self.discriminator,
self.g_params['z_dim'], labels=condition)
fake_images, 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)
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
cond_loss = tf.reduce_sum(cond_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 + cond_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, cond_loss
return d_loss, d_gan_loss, r1_penalty, e_loss, c_loss
def g_train_step(self, dist_inputs):
real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
real_images, conditions, mesh_static, mesh_body_posed, mask = 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'], labels=condition)
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)
# 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
g_loss = g_loss + e_loss + c_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
return g_loss, e_loss, c_loss
def g_train_step_reg(self, dist_inputs):
real_images, condition = dist_inputs[0][0], dist_inputs[0][1]
real_images, conditions, mesh_static, mesh_body_posed, mask = 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,
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, labels=condition)
self.pl_minibatch_shrink, self.pl_denorm, self.pl_decay, conditions=conditions)
pl_penalty = pl_penalty * self.pl_weight * self.g_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 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
# combine
g_loss = g_gan_loss + pl_penalty
g_loss = g_gan_loss + pl_penalty + e_loss + c_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
return g_loss, g_gan_loss, pl_penalty, e_loss, c_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[0], axis=None)
mean_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[1], axis=None)
return mean_d_loss, mean_cond_loss
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):
@@ -200,14 +235,17 @@ 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_cond_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
return mean_d_loss, mean_d_gan_loss, mean_r1_penalty, mean_cond_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):
per_replica_losses = strategy.run(fn=self.g_train_step, args=(inputs,))
mean_g_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
return mean_g_loss
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
@tf.function
def dist_g_train_step_reg(inputs):
@@ -215,7 +253,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)
return mean_g_loss, mean_g_gan_loss, mean_pl_penalty
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
@tf.function
def dist_gen_samples(dist_inputs):
@@ -238,28 +278,29 @@ class Trainer(object):
metric_g_gan_loss = tf.keras.metrics.Mean('g_gan_loss', dtype=tf.float32)
metric_r1_penalty = tf.keras.metrics.Mean('r1_penalty', dtype=tf.float32)
metric_pl_penalty = tf.keras.metrics.Mean('pl_penalty', dtype=tf.float32)
metric_cond_loss = tf.keras.metrics.Mean('cond_loss', 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)
# start training
zero = tf.constant(0.0, dtype=tf.float32)
print('max_steps: {}'.format(self.max_steps))
t_start = time.time()
step = self.g_optimizer.iterations.numpy()
for real_images in dist_datasets:
for batch in dist_datasets:
# d train step
if (step + 1) % self.d_opt['reg_interval'] == 0:
d_loss, d_gan_loss, r1_penalty, cond_loss = dist_d_train_step_reg((real_images, ))
d_loss, d_gan_loss, r1_penalty, e_loss, c_loss = dist_d_train_step_reg((batch, ))
else:
d_loss, cond_loss = dist_d_train_step((real_images, ))
d_loss, e_loss, c_loss = dist_d_train_step((batch, ))
d_gan_loss = d_loss
r1_penalty = zero
# g train step
if (step + 1) % self.g_opt['reg_interval'] == 0:
g_loss, g_gan_loss, pl_penalty = dist_g_train_step_reg((real_images, ))
g_loss, g_gan_loss, pl_penalty, e_loss, c_loss = dist_g_train_step_reg((batch, ))
else:
g_loss = dist_g_train_step((real_images,))
g_loss, e_loss, c_loss = dist_g_train_step((batch,))
g_gan_loss = g_loss
pl_penalty = zero
@@ -273,7 +314,8 @@ class Trainer(object):
metric_g_gan_loss(g_gan_loss)
metric_r1_penalty(r1_penalty)
metric_pl_penalty(pl_penalty)
metric_cond_loss(cond_loss)
metric_e_loss(e_loss)
metric_c_loss(c_loss)
# get current step
step += 1
@@ -286,7 +328,8 @@ class Trainer(object):
tf.summary.scalar('g_gan_loss', metric_g_gan_loss.result(), step=step)
tf.summary.scalar('r1_penalty', metric_r1_penalty.result(), step=step)
tf.summary.scalar('pl_penalty', metric_pl_penalty.result(), step=step)
tf.summary.scalar('cond_loss', metric_cond_loss.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)
# save every self.save_step
if step % self.save_step == 0:
@@ -296,8 +339,8 @@ 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_labels = tf.ones((self.n_samples, self.g_params['labels_dim']), dtype=tf.dtypes.float32)
summary_image = dist_gen_samples((test_z, test_labels))
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))
# convert to tensor image
summary_image = self.convert_per_replica_image(summary_image, strategy)
@@ -312,7 +355,8 @@ 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()))
r1_penalty.numpy(), pl_penalty.numpy(),
e_loss.numpy(), c_loss.numpy()))
# reset timer
t_start = time.time()
@@ -364,7 +408,7 @@ def main():
parser.add_argument('--tfrecord_dir', default='./tfrecords', type=str)
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=12, type=int)
parser.add_argument('--batch_size_per_replica', default=4, type=int)
args = vars(parser.parse_args())
# check tensorflow version
@@ -382,15 +426,15 @@ def main():
# featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res'])
g_params = {
'z_dim': 512,
'w_dim': 512,
'labels_dim': 83,
'z_dim': 256,
'w_dim': 256,
'labels_dim': -1,
'n_mapping': 8,
'resolutions': train_resolutions,
'featuremaps': train_featuremaps,
}
d_params = {
'labels_dim': 83,
'labels_dim': -1,
'resolutions': train_resolutions,
'featuremaps': train_featuremaps,
}
@@ -401,7 +445,14 @@ def main():
# prepare dataset
# dataset = get_ffhq_dataset(args['tfrecord_dir'], args['train_res'], batch_size=global_batch_size, epochs=None)
dataset = DatasetSimple('/workspace/Dataset/stylegan').batch(global_batch_size).shuffle(buffer_size=1000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE).cache().repeat()
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
dataset = dataset.with_options(options)