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

add posed mesh, metrics, tanh, l1 with gt

Browse Source
This commit is contained in:
luni
2022-06-17 19:41:16 +00:00
parent cf1c536260
commit a14acceeb6
5 changed files with 87 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+3 -2
View File
@@ -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, )
)
stylegan2/layers/bias_act.py
+4 -1
View File
@@ -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)
stylegan2/layers/synthesis_block.py
+7 -2
View File
@@ -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
stylegan2/layers/to_rgb.py
+26
View File
@@ -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
train.py
+47 -28
View File
@@ -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