diff --git a/dataset.py b/dataset.py index 750ffba..42bc964 100644 --- a/dataset.py +++ b/dataset.py @@ -26,14 +26,15 @@ class DatasetSimple(tf.data.Dataset): mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32) uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32) + mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask - yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask + yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed def __new__(cls, folder, res=256): return tf.data.Dataset.from_generator( cls._generator, - output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32) ), + output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), tf.TensorSpec(shape = (14475,), dtype = tf.float32), tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32) ), args=(folder, ) ) diff --git a/stylegan2/layers/bias_act.py b/stylegan2/layers/bias_act.py index bea2284..afff595 100644 --- a/stylegan2/layers/bias_act.py +++ b/stylegan2/layers/bias_act.py @@ -7,12 +7,15 @@ import tensorflow as tf class BiasAct(tf.keras.layers.Layer): def __init__(self, lrmul, act, **kwargs): super(BiasAct, self).__init__(**kwargs) - assert act in ['linear', 'lrelu'] + assert act in ['linear', 'lrelu', 'tanh'] self.lrmul = lrmul if act == 'linear': self.act = tf.keras.layers.Lambda(lambda x: tf.identity(x)) self.gain = 1.0 + elif act == 'tanh': + self.act = tf.keras.layers.Activation('tanh') + self.gain = 1.0 else: self.act = tf.keras.layers.LeakyReLU(alpha=0.2) self.gain = np.sqrt(2) diff --git a/stylegan2/layers/synthesis_block.py b/stylegan2/layers/synthesis_block.py index 9edf2a0..b4f8c45 100644 --- a/stylegan2/layers/synthesis_block.py +++ b/stylegan2/layers/synthesis_block.py @@ -3,7 +3,7 @@ import tensorflow as tf from stylegan2.layers.modulated_conv2d import ModulatedConv2D from stylegan2.layers.bias_act import BiasAct from stylegan2.layers.noise import Noise -from stylegan2.layers.to_rgb import ToRGB +from stylegan2.layers.to_rgb import ToRGB, ToRGBLast from stylegan2.layers.cuda.upfirdn_2d_v2 import upsample_2d, compute_paddings @@ -120,12 +120,17 @@ class Synthesis(tf.keras.layers.Layer): prev_n_f = n_f self.blocks = list() self.torgbs = list() - for res, n_f in zip(self.resolutions[1:], self.featuremaps[1:]): + for res, n_f in zip(self.resolutions[1:-1], self.featuremaps[1:-1]): self.blocks.append(SynthesisBlock(in_ch=prev_n_f, fmaps=n_f, res=res, name='{:d}x{:d}/block'.format(res, res))) self.torgbs.append(ToRGB(in_ch=n_f, res=res * 2, name='{:d}x{:d}/ToRGB'.format(res, res))) prev_n_f = n_f + res, n_f = self.resolutions[-1], self.featuremaps[-1] + self.blocks.append(SynthesisBlock(in_ch=prev_n_f, fmaps=n_f, res=res, + name='{:d}x{:d}/block'.format(res, res))) + self.torgbs.append(ToRGBLast(in_ch=n_f, res=res * 2, name='{:d}x{:d}/ToRGB'.format(res, res))) + def call(self, inputs, training=None, mask=None): w_broadcasted = inputs diff --git a/stylegan2/layers/to_rgb.py b/stylegan2/layers/to_rgb.py index 78a2121..76e8403 100644 --- a/stylegan2/layers/to_rgb.py +++ b/stylegan2/layers/to_rgb.py @@ -28,3 +28,29 @@ class ToRGB(tf.keras.layers.Layer): 'res': self.res, }) return config + + +class ToRGBLast(tf.keras.layers.Layer): + def __init__(self, in_ch, res, **kwargs): + super(ToRGBLast, self).__init__(**kwargs) + self.in_ch = in_ch + self.res = res + + self.conv = ModulatedConv2D(in_res=res, in_fmaps=in_ch, fmaps=3, kernel=1, up=False, down=False, demodulate=False, + resample_kernel=None, gain=1.0, lrmul=1.0, fused_modconv=True, name='conv') + self.apply_bias = BiasAct(lrmul=1.0, act='tanh', name='bias') + + def call(self, inputs, training=None, mask=None): + x, w = inputs + + x = self.conv([x, w]) + x = self.apply_bias(x) + return x + + def get_config(self): + config = super(ToRGBLast, self).get_config() + config.update({ + 'in_ch': self.in_ch, + 'res': self.res, + }) + return config \ No newline at end of file diff --git a/train.py b/train.py index 416ddfa..cc535a3 100644 --- a/train.py +++ b/train.py @@ -17,8 +17,8 @@ from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_l from losses import collision_loss as collision_penalty -os.environ["CUDA_VISIBLE_DEVICES"] = "3" -run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_19', notes='learning offsets, all losses from the start, smaller generator (from 128), slower learning rate for generator, sgd for discriminator') +os.environ["CUDA_VISIBLE_DEVICES"] = "4" +run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_23', 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, trained from test-22-10000') artifact = wandb.Artifact('code', type='code') artifact.add_file('train.py') artifact.add_file('losses.py') @@ -128,7 +128,7 @@ 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 @@ -144,7 +144,7 @@ class Trainer(object): 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 @@ -165,7 +165,7 @@ class Trainer(object): def g_train_step(self, dist_inputs): - real_images, conditions, _, _, _ = dist_inputs[0] + real_images, conditions, _, _, _, _ = dist_inputs[0] with tf.GradientTape() as g_tape: # compute losses @@ -179,7 +179,7 @@ class Trainer(object): return g_loss def g_train_step_reg(self, dist_inputs): - real_images, conditions, _, _, _ = dist_inputs[0] + real_images, conditions, _, _, _, _ = dist_inputs[0] with tf.GradientTape() as g_tape: # compute losses @@ -201,34 +201,37 @@ class Trainer(object): def g_train_step_extra(self, dist_inputs): - real_images, conditions, mesh_static, mesh_body_posed, mesh_mask = dist_inputs[0] - uv_mask = conditions[4] + real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed = dist_inputs[0] + uv_mask = conditions[3] 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, mesh_mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) - l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) + e_loss = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask) + c_loss = 1e5 * 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) + 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 + 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] - real_images, conditions, mesh_static, mesh_body_posed, mesh_mask = dist_inputs[0] - uv_mask = conditions[4] + real_images, conditions, mesh_static, mesh_body_posed, mesh_mask, mesh_posed = dist_inputs[0] + uv_mask = conditions[3] with tf.GradientTape() as g_tape: # compute losses @@ -239,9 +242,11 @@ class Trainer(object): mesh_pred = uv2mesh(fake_images) - e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask) - c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) - l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask) + e_loss = 1e4 * edge_loss(mesh_pred, mesh_static, mesh_mask) + c_loss = 1e5 * 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) + 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 @@ -249,13 +254,14 @@ class Trainer(object): 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_loss + self.unsupervised_coeff * (e_loss + c_loss) + self.gt_loss_coeff * l1_loss + 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 + return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss, euclid_error, ms_ssim def train(self, dist_datasets, strategy): @tf.function @@ -293,7 +299,9 @@ class Trainer(object): 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) - return mean_g_loss, mean_e_loss, mean_c_loss, mean_l1_loss + 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): @@ -304,7 +312,9 @@ class Trainer(object): 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) - return mean_g_loss, mean_g_gan_loss, mean_pl_penalty, mean_e_loss, mean_c_loss, mean_l1_loss + 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): @@ -332,6 +342,8 @@ class Trainer(object): 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) @@ -351,20 +363,24 @@ class Trainer(object): # 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 = dist_g_train_step_reg_extra((batch, )) + 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 = dist_g_train_step_extra((batch,)) + 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 @@ -374,7 +390,7 @@ class Trainer(object): # self.unsupervised_coeff = 1 if step < self.disable_l1_loss: - self.gt_loss_coeff = 1 - 1/step + self.gt_loss_coeff = 1 - step / self.disable_l1_loss else: self.gt_loss_coeff = 0 @@ -391,6 +407,8 @@ class Trainer(object): 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) # get current step step += 1 @@ -408,6 +426,8 @@ class Trainer(object): 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: @@ -428,9 +448,8 @@ class Trainer(object): tf.summary.image('images', summary_image, step=step) template = uv2mesh(batch[1][2]) - mesh_gt = uv2mesh(batch[0]) # gt = uv2mesh(batch[0]) - plot_meshes(mesh_writer, garment=mesh, mesh_gt=mesh_gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step) + plot_meshes(mesh_writer, garment=mesh, mesh_gt=batch[5], 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() @@ -584,7 +603,7 @@ def main(): 'd_params': d_params, # training params - 'g_opt': {'learning_rate': 0.0002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8}, + 'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8}, '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, @@ -592,7 +611,7 @@ def main(): 'train_res': args['train_res'], } - trainer = Trainer(training_parameters, name=f'stylegan2-test-19-{args["train_res"]}x{args["train_res"]}') + trainer = Trainer(training_parameters, name=f'stylegan2-test-23-{args["train_res"]}x{args["train_res"]}') trainer.train(dist_dataset, strategy) return