add posed mesh, metrics, tanh, l1 with gt

2022-06-17 19:41:16 +00:00
parent cf1c536260
commit a14acceeb6
5 changed files with 87 additions and 33 deletions
@@ -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