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

test-20

Browse Source
This commit is contained in:
luni
2022-06-17 14:30:35 +00:00
parent e712e456bc
commit cf1c536260
3 changed files with 101 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
custom_utils/custom_utils.py
+51 -1
View File
@@ -1,5 +1,6 @@
import numpy as np
import tensorflow as tf
import plotly.graph_objects as go
def m2uv_array(m2uv):
@@ -85,4 +86,53 @@ def maskFacesBatch(V, F, mask):
v, f = _maskFaces(V[i], F, mask[i])
Vs.append(v)
Fs.append(f)
return Vs, Fs
return Vs, Fs
def plot_3d(mesh_gt, mesh_pred, faces, collision, collision_faces, collision_color):
mesh_pred[:,1] += 1
collision[:,1] += 2
fig = go.Figure(data=[
go.Mesh3d(
# 8 vertices of a cube
x=mesh_gt[:,0],
y=mesh_gt[:,1],
z=mesh_gt[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2]
),
go.Mesh3d(
# 8 vertices of a cube
x=mesh_pred[:,0],
y=mesh_pred[:,1],
z=mesh_pred[:,2],
# i, j and k give the vertices of triangles
i = faces[:,0],
j = faces[:,1],
k = faces[:,2],
),
go.Mesh3d(
# 8 vertices of a cube
x=collision[:,0],
y=collision[:,1],
z=collision[:,2],
# i, j and k give the vertices of triangles
i = collision_faces[:,0],
j = collision_faces[:,1],
k = collision_faces[:,2],
vertexcolor=collision_color
)
])
fig.update_layout(scene_aspectmode='manual',
scene_aspectratio=dict(x=1, y=2, z=1))
fig.update_layout(
scene = dict(
xaxis = dict(nticks=4, range=[-1,1],),
yaxis = dict(nticks=4, range=[-1,3],),
zaxis = dict(nticks=4, range=[-1,1],),)
)
return fig
requierements.txt
+2 -1
View File
@@ -4,4 +4,5 @@ opencv-python
wandb
OpenEXR
tensorflow-graphics
scipy
scipy
plotly
train.py
+48 -29
View File
@@ -5,7 +5,7 @@ import wandb
from cv2 import drawFrameAxes
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, plot_3d
from tensorboardX import SummaryWriter
from utils import str_to_bool
@@ -18,7 +18,7 @@ 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_13(test-16)', notes='learning offsets, all losses from the start')
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')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('train.py')
artifact.add_file('losses.py')
@@ -46,6 +46,7 @@ class Trainer(object):
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']
@@ -53,6 +54,7 @@ class Trainer(object):
self.save_step = 5000
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}',
@@ -86,6 +88,7 @@ class Trainer(object):
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'],
@@ -199,6 +202,7 @@ 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]
with tf.GradientTape() as g_tape:
# compute losses
@@ -207,23 +211,24 @@ 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)
m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask)
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
# 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
m_loss = tf.reduce_sum(m_loss) * self.global_batch_scaler
l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
g_loss = g_loss + self.unsupervised_coeff * (e_loss + c_loss + m_loss)
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, m_loss
return g_loss, e_loss, c_loss, l1_loss
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]
with tf.GradientTape() as g_tape:
# compute losses
@@ -236,21 +241,21 @@ class Trainer(object):
e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask)
c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask)
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(real_images - fake_images)), axis=1) * uv_mask)
# 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
m_loss = tf.reduce_sum(m_loss) * self.global_batch_scaler
l1_loss = tf.reduce_sum(l1_loss) * self.global_batch_scaler
# combine
g_loss = g_gan_loss + pl_penalty + self.unsupervised_coeff * (e_loss + c_loss + m_loss)
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, g_gan_loss, pl_penalty, e_loss, c_loss, m_loss
return g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss
def train(self, dist_datasets, strategy):
@tf.function
@@ -287,8 +292,8 @@ class Trainer(object):
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_m_loss = strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses[3], axis=None)
return mean_g_loss, mean_e_loss, mean_c_loss, mean_m_loss
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
@tf.function
def dist_g_train_step_reg_extra(inputs):
@@ -298,8 +303,8 @@ class Trainer(object):
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_m_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_m_loss
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
@tf.function
def dist_gen_samples(dist_inputs):
@@ -326,7 +331,7 @@ class Trainer(object):
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_m_loss = tf.keras.metrics.Mean('uv_mask_loss', dtype=tf.float32)
metric_l1_loss = tf.keras.metrics.Mean('l1_loss', dtype=tf.float32)
# start training
zero = tf.constant(0.0, dtype=tf.float32)
@@ -346,20 +351,20 @@ 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, m_loss = dist_g_train_step_reg_extra((batch, ))
g_loss, g_gan_loss, pl_penalty, e_loss, c_loss, l1_loss = 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
m_loss = 0
l1_loss = 0
else:
if step > self.enable_unsupervised_losses:
g_loss, e_loss, c_loss, m_loss = dist_g_train_step_extra((batch,))
g_loss, e_loss, c_loss, l1_loss = dist_g_train_step_extra((batch,))
else:
g_loss = dist_g_train_step((batch,))
e_loss = 0
c_loss = 0
m_loss = 0
l1_loss = 0
g_gan_loss = g_loss
pl_penalty = zero
@@ -367,6 +372,11 @@ class Trainer(object):
# if step > self.enable_unsupervised_losses:
# self.unsupervised_coeff = 1
if step < self.disable_l1_loss:
self.gt_loss_coeff = 1 - 1/step
else:
self.gt_loss_coeff = 0
# update g_clone
self.g_clone.set_as_moving_average_of(self.generator)
@@ -380,7 +390,7 @@ class Trainer(object):
metric_pl_penalty(pl_penalty)
metric_e_loss(e_loss)
metric_c_loss(c_loss)
metric_m_loss(m_loss)
metric_l1_loss(l1_loss)
# get current step
step += 1
@@ -396,7 +406,8 @@ class Trainer(object):
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('uv_mask_loss', metric_m_loss.result(), 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)
# save every self.save_step
if step % self.save_step == 0:
@@ -417,8 +428,9 @@ 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, template=template, body=batch[3], mask=batch[4].numpy(), step=step)
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=gt, template=template, body=batch[3], mask=batch[4].numpy(), step=step)
# train_summary_writer.flush()
@@ -472,11 +484,14 @@ def filter_resolutions_featuremaps(resolutions, featuremaps, res):
return filtered_resolutions, filtered_featuremaps
def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0):
def plot_meshes(writer, garment, mesh_gt, template, body, mask, step, nr_offset=0):
body = body.numpy()
template = template.numpy()
mesh_gt = mesh_gt.numpy()
F = np.load("custom_utils/smpl_faces.npy")
garment, faces = maskFacesBatch(garment.numpy(), F, mask.astype(bool))
template, _ = maskFacesBatch(template.numpy(), F, mask.astype(bool))
body = body.numpy()
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))]
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))]
@@ -488,7 +503,10 @@ def plot_meshes(writer, garment, template, body, mask, step, nr_offset=0):
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[i], faces[i], collision[i], collision_faces[i], collision_color[i])
wandb.log({"garment_" + str(i): plotly_fig})
def main():
# global program arguments parser
@@ -515,7 +533,7 @@ def main():
# network params
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
g_featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16]
g_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]
@@ -559,13 +577,14 @@ def main():
'use_custom_cuda': args['use_custom_cuda'],
'model_base_dir': args['model_base_dir'],
'enable_unsupervised_losses': 0,
'disable_l1_loss': 10000,
# network params
'g_params': g_params,
'd_params': d_params,
# training params
'g_opt': {'learning_rate': 0.002, 'beta1': 0.0, 'beta2': 0.99, 'epsilon': 1e-08, 'reg_interval': 8},
'g_opt': {'learning_rate': 0.0002, '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,
@@ -573,7 +592,7 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-16-{args["train_res"]}x{args["train_res"]}')
trainer = Trainer(training_parameters, name=f'stylegan2-test-19-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, strategy)
return