662 lines
34 KiB
Python
662 lines
34 KiB
Python
import os
|
|
import time
|
|
import argparse
|
|
import wandb
|
|
from cv2 import drawFrameAxes
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
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, 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"] = "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')
|
|
artifact.add_file('stylegan2/layers/mapping_block.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)
|
|
generator = load_generator(g_params=g_params, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
|
g_clone = load_generator(g_params=g_params, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
|
|
|
# set initial g_clone weights same as generator
|
|
g_clone.set_weights(generator.get_weights())
|
|
return discriminator, generator, g_clone
|
|
|
|
|
|
uv2mesh = UV2Mesh(res=256)
|
|
smpl = SMPL()
|
|
|
|
class Trainer(object):
|
|
def __init__(self, t_params, name):
|
|
# 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']
|
|
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.disable_l1_loss = t_params['disable_l1_loss']
|
|
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 = 100
|
|
self.save_step = 10000
|
|
self.image_summary_step = 500
|
|
self.unsupervised_coeff = 1
|
|
self.gt_loss_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}')
|
|
|
|
# copy network params
|
|
self.g_params = t_params['g_params']
|
|
self.d_params = t_params['d_params']
|
|
|
|
# set optimizer params
|
|
self.global_batch_scaler = 1.0 / float(self.global_batch_size)
|
|
self.r1_gamma = 10.0
|
|
self.g_opt = self.update_optimizer_params(t_params['g_opt'])
|
|
self.d_opt = self.update_optimizer_params(t_params['d_opt'])
|
|
self.pl_minibatch_shrink = 2
|
|
self.pl_weight = float(self.pl_minibatch_shrink)
|
|
self.pl_denorm = tf.math.rsqrt(float(self.train_res) * float(self.train_res))
|
|
self.pl_decay = 0.01
|
|
self.pl_mean = tf.Variable(initial_value=0.0, name='pl_mean', trainable=False,
|
|
synchronization=tf.VariableSynchronization.ON_READ,
|
|
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
|
|
|
|
# create model: model and optimizer must be created under `strategy.scope`
|
|
self.discriminator, self.generator, self.g_clone = initiate_models(self.g_params,
|
|
self.d_params,
|
|
self.use_custom_cuda)
|
|
|
|
# set optimizers
|
|
self.d_optimizer = tf.keras.optimizers.Adam(self.d_opt['learning_rate'],
|
|
beta_1=self.d_opt['beta1'],
|
|
beta_2=self.d_opt['beta2'],
|
|
epsilon=self.d_opt['epsilon'])
|
|
# self.d_optimizer = tf.keras.optimizers.SGD()
|
|
self.g_optimizer = tf.keras.optimizers.Adam(self.g_opt['learning_rate'],
|
|
beta_1=self.g_opt['beta1'],
|
|
beta_2=self.g_opt['beta2'],
|
|
epsilon=self.g_opt['epsilon'])
|
|
|
|
# setup saving locations (object based savings)
|
|
self.ckpt_dir = os.path.join(self.model_base_dir, name)
|
|
self.ckpt = tf.train.Checkpoint(d_optimizer=self.d_optimizer,
|
|
g_optimizer=self.g_optimizer,
|
|
discriminator=self.discriminator,
|
|
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=20)
|
|
|
|
# try to restore
|
|
self.ckpt.restore(self.manager.latest_checkpoint)
|
|
if self.manager.latest_checkpoint:
|
|
print('Restored from {}'.format(self.manager.latest_checkpoint))
|
|
|
|
# check if already trained in this resolution
|
|
restored_step = self.g_optimizer.iterations.numpy()
|
|
if restored_step >= self.max_steps:
|
|
print('Already reached max steps {}/{}'.format(restored_step, self.max_steps))
|
|
self.reached_max_steps = True
|
|
return
|
|
else:
|
|
print('Not restoring from saved checkpoint')
|
|
|
|
@staticmethod
|
|
def update_optimizer_params(params):
|
|
params_copy = params.copy()
|
|
mb_ratio = params_copy['reg_interval'] / (params_copy['reg_interval'] + 1)
|
|
params_copy['learning_rate'] = params_copy['learning_rate'] * mb_ratio
|
|
params_copy['beta1'] = params_copy['beta1'] ** mb_ratio
|
|
params_copy['beta2'] = params_copy['beta2'] ** mb_ratio
|
|
return params_copy
|
|
|
|
def d_train_step(self, dist_inputs):
|
|
# 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=batch['conditions'])
|
|
|
|
# scale loss
|
|
d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
|
|
|
|
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
|
|
|
|
def d_train_step_reg(self, dist_inputs):
|
|
# 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=batch['conditions'])
|
|
r1_penalty = r1_penalty * (0.5 * self.r1_gamma) * self.d_opt['reg_interval']
|
|
|
|
# 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
|
|
|
|
# combine
|
|
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
|
|
|
|
|
|
def g_train_step(self, dist_inputs):
|
|
# 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=batch['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):
|
|
# 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=batch['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):
|
|
# 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=batch['conditions'])
|
|
|
|
mesh_pred_unposed = uv2mesh(fake_images)
|
|
mesh_pred = smpl.run_garment(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed)
|
|
|
|
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
|
|
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
|
|
|
|
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, l1_loss, euclid_error, ms_ssim
|
|
|
|
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]
|
|
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=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(batch['conditions']['pose_shape'][:, :72], batch['conditions']['pose_shape'][:, 72:], mesh_pred_unposed)
|
|
|
|
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
|
|
ms_ssim = tf.reduce_sum(ms_ssim) * self.global_batch_scaler
|
|
|
|
# combine
|
|
g_loss = g_gan_loss + pl_penalty + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_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, l1_loss, euclid_error, ms_ssim
|
|
|
|
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,))
|
|
mean_d_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
|
|
return mean_d_loss
|
|
|
|
@tf.function
|
|
def dist_d_train_step_reg(inputs):
|
|
per_replica_losses = strategy.run(fn=self.d_train_step_reg, args=(inputs,))
|
|
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)
|
|
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, axis=None)
|
|
return mean_g_loss
|
|
|
|
@tf.function
|
|
def dist_g_train_step_reg(inputs):
|
|
per_replica_losses = strategy.run(fn=self.g_train_step_reg, 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)
|
|
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_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
|
|
|
|
@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)
|
|
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
|
|
|
|
@tf.function
|
|
def dist_gen_samples(dist_inputs):
|
|
per_replica_samples = strategy.run(self.gen_samples, args=(dist_inputs,))
|
|
return per_replica_samples
|
|
|
|
if self.reached_max_steps:
|
|
return
|
|
|
|
# start actual training
|
|
print('Start Training')
|
|
|
|
# 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)
|
|
metric_g_loss = tf.keras.metrics.Mean('g_loss', dtype=tf.float32)
|
|
metric_d_gan_loss = tf.keras.metrics.Mean('d_gan_loss', dtype=tf.float32)
|
|
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_e_loss = tf.keras.metrics.Mean('edge_loss', dtype=tf.float32)
|
|
metric_c_loss = tf.keras.metrics.Mean('collision_loss', dtype=tf.float32)
|
|
metric_l1_loss = tf.keras.metrics.Mean('l1_loss', dtype=tf.float32)
|
|
metric_euclid_error = tf.keras.metrics.Mean('euclid_error', dtype=tf.float32)
|
|
metric_ms_ssim = tf.keras.metrics.Mean('ms_ssim', 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 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:
|
|
d_loss, d_gan_loss, r1_penalty = dist_d_train_step_reg((batch, ))
|
|
else:
|
|
d_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:
|
|
if step > self.enable_unsupervised_losses:
|
|
g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss, euclid_error, ms_ssim = 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
|
|
l1_loss = 0
|
|
euclid_error = 0
|
|
ms_ssim = 0
|
|
else:
|
|
if step > self.enable_unsupervised_losses:
|
|
g_loss, e_loss, c_loss, l1_loss, euclid_error, ms_ssim = dist_g_train_step_extra((batch,))
|
|
else:
|
|
g_loss = dist_g_train_step((batch,))
|
|
e_loss = 0
|
|
c_loss = 0
|
|
l1_loss = 0
|
|
euclid_error = 0
|
|
ms_ssim = 0
|
|
|
|
g_gan_loss = g_loss
|
|
pl_penalty = zero
|
|
|
|
|
|
# if step > self.enable_unsupervised_losses:
|
|
# self.unsupervised_coeff = 1
|
|
|
|
if step < self.disable_l1_loss:
|
|
self.gt_loss_coeff = 1 - step / self.disable_l1_loss
|
|
else:
|
|
self.gt_loss_coeff = 0
|
|
|
|
# update g_clone
|
|
self.g_clone.set_as_moving_average_of(self.generator)
|
|
|
|
# update metrics
|
|
metric_d_loss(d_loss)
|
|
metric_g_loss(g_loss)
|
|
metric_d_gan_loss(d_gan_loss)
|
|
metric_g_gan_loss(g_gan_loss)
|
|
metric_r1_penalty(r1_penalty)
|
|
metric_pl_penalty(pl_penalty)
|
|
metric_e_loss(e_loss)
|
|
metric_c_loss(c_loss)
|
|
metric_l1_loss(l1_loss)
|
|
metric_euclid_error(euclid_error)
|
|
metric_ms_ssim(ms_ssim)
|
|
|
|
# save to tensorboard
|
|
with train_summary_writer.as_default():
|
|
tf.summary.scalar('d_loss', metric_d_loss.result(), step=step)
|
|
tf.summary.scalar('g_loss', metric_g_loss.result(), step=step)
|
|
tf.summary.scalar('d_gan_loss', metric_d_gan_loss.result(), step=step)
|
|
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('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('gt_loss_coeff', self.gt_loss_coeff, step=step)
|
|
tf.summary.scalar('l1_loss', metric_l1_loss.result(), step=step)
|
|
tf.summary.scalar('euclid_error', metric_euclid_error.result(), step=step)
|
|
tf.summary.scalar('ms_ssim', metric_ms_ssim.result(), step=step)
|
|
|
|
# save every self.save_step
|
|
if step % self.save_step == 0:
|
|
self.manager.save(checkpoint_number=step)
|
|
|
|
# save every self.image_summary_step
|
|
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['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['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)
|
|
|
|
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()
|
|
|
|
# print every self.print_steps
|
|
if step % self.print_step == 0:
|
|
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()))
|
|
|
|
# reset timer
|
|
t_start = time.time()
|
|
|
|
# check exit status
|
|
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)
|
|
return
|
|
|
|
def gen_samples(self, inputs):
|
|
test_z, test_conditions = inputs
|
|
|
|
# run networks
|
|
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)
|
|
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
|
|
def convert_per_replica_image(nchw_per_replica_images, strategy):
|
|
as_tensor = tf.concat(strategy.experimental_local_results(nchw_per_replica_images), axis=0)
|
|
as_tensor = tf.transpose(as_tensor, perm=[0, 2, 3, 1])
|
|
as_tensor = (tf.clip_by_value(as_tensor, -1.0, 1.0) + 1.0) * 127.5
|
|
as_tensor = tf.cast(as_tensor, tf.uint8)
|
|
return as_tensor
|
|
|
|
|
|
def filter_resolutions_featuremaps(resolutions, featuremaps, res):
|
|
index = resolutions.index(res)
|
|
filtered_resolutions = resolutions[:index + 1]
|
|
filtered_featuremaps = featuremaps[:index + 1]
|
|
return filtered_resolutions, filtered_featuremaps
|
|
|
|
|
|
def plot_meshes(writer, garment_unposed, garment, mesh_gt, template, body, mask, step, nr_offset=0):
|
|
body = body.numpy()
|
|
template = template.numpy()
|
|
mesh_gt = mesh_gt.numpy()
|
|
garment_unposed = garment_unposed.numpy()
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
garment, faces = maskFacesBatch(garment.numpy(), F, mask.astype(bool))
|
|
template = [template[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
|
|
mesh_gt = [mesh_gt[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
|
|
garment_unposed = [garment_unposed[i, mask.astype(bool)[i,:], :] for i in range(len(garment))]
|
|
|
|
collision = [np.concatenate((body[i], garment[i]), 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[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])
|
|
wandb.log({"garment_" + str(i): plotly_fig})
|
|
|
|
def main():
|
|
# global program arguments parser
|
|
parser = argparse.ArgumentParser(description='')
|
|
parser.add_argument('--allow_memory_growth', type=str_to_bool, nargs='?', const=True, default=True)
|
|
parser.add_argument('--debug_split_gpu', type=str_to_bool, nargs='?', const=True, default=False)
|
|
parser.add_argument('--use_tf_function', type=str_to_bool, nargs='?', const=True, default=True)
|
|
parser.add_argument('--use_custom_cuda', type=str_to_bool, nargs='?', const=True, default=True)
|
|
parser.add_argument('--model_base_dir', default='./models', type=str)
|
|
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=4, type=int)
|
|
parser.add_argument('--name', type=str)
|
|
args = vars(parser.parse_args())
|
|
|
|
# check tensorflow version
|
|
# cur_tf_ver = check_tf_version()
|
|
|
|
# GPU environment settings
|
|
if args['allow_memory_growth']:
|
|
allow_memory_growth()
|
|
if args['debug_split_gpu']:
|
|
split_gpu_for_testing(mem_in_gb=4.5)
|
|
|
|
# network params
|
|
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
|
|
g_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8]
|
|
d_featuremaps = [128, 128, 128, 128, 128, 64, 32, 16, 8]
|
|
# d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
|
|
# featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
|
|
# featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
|
|
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': 128,
|
|
'w_dim': 128,
|
|
'labels_dim': -1,
|
|
'n_mapping': 8,
|
|
'resolutions': train_resolutions,
|
|
'featuremaps': train_g_featuremaps,
|
|
'zero_body': np.load('custom_utils/uv_zero_body.npy')
|
|
}
|
|
d_params = {
|
|
'batch_size': args['batch_size_per_replica'],
|
|
'labels_dim': -1,
|
|
'resolutions': train_resolutions,
|
|
'featuremaps': train_d_featuremaps,
|
|
}
|
|
|
|
# prepare distribute strategy
|
|
strategy = tf.distribute.MirroredStrategy()
|
|
global_batch_size = args['batch_size_per_replica'] * strategy.num_replicas_in_sync
|
|
|
|
# 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()
|
|
|
|
factory = DatasetFactory()
|
|
signature_dynamic = { 'id': (), 'frame_nr': (),
|
|
'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
|
|
dataset = dataset.with_options(options)
|
|
|
|
with strategy.scope():
|
|
# distribute dataset
|
|
dist_dataset = strategy.experimental_distribute_dataset(dataset)
|
|
|
|
# training parameters
|
|
training_parameters = {
|
|
# global params
|
|
# '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'],
|
|
'enable_unsupervised_losses': 0,
|
|
'disable_l1_loss': 30000,
|
|
|
|
# network params
|
|
'g_params': g_params,
|
|
'd_params': d_params,
|
|
|
|
# training params
|
|
'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,
|
|
'n_samples': 12,
|
|
'train_res': 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
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|