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

train on template only with mask applied

Browse Source
This commit is contained in:
luni
2022-06-14 10:16:55 +00:00
parent def8a85cf9
commit 245f811d47
11 changed files with 112 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
train.py
+21 -20
View File
@@ -17,12 +17,13 @@ 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"] = "2"
run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_7(test-10)')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('train.py')
artifact.add_file('losses.py')
run.log_artifact(artifact)
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
# run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_10(test-13)')
# 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)
@@ -124,13 +125,12 @@ 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
d_loss = d_logistic(real_images, self.generator, self.discriminator, self.g_params['z_dim'], conditions=conditions)
# scale loss
d_loss = tf.reduce_sum(d_loss) * self.global_batch_scaler
@@ -141,7 +141,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
@@ -162,7 +162,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
@@ -176,7 +176,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
@@ -198,16 +198,16 @@ class Trainer(object):
def g_train_step_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]
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 = 0.1 * edge_loss(mesh_pred, mesh_static, mesh_mask)
c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
m_loss = uv_mask_loss(fake_images, uv_mask)
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)
# scale loss
g_loss = tf.reduce_sum(g_loss) * self.global_batch_scaler
@@ -222,7 +222,8 @@ class Trainer(object):
return g_loss, e_loss, c_loss, m_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, uv_mask = dist_inputs[0]
real_images, conditions, mesh_static, mesh_body_posed, mesh_mask = dist_inputs[0]
with tf.GradientTape() as g_tape:
# compute losses
@@ -233,9 +234,9 @@ class Trainer(object):
mesh_pred = uv2mesh(fake_images)
e_loss = 0.1 * edge_loss(mesh_pred, mesh_static, mesh_mask)
c_loss = 0.0 # collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
m_loss = uv_mask_loss(fake_images, uv_mask)
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)
# scale loss
pl_penalty = tf.reduce_sum(pl_penalty) * self.global_batch_scaler
@@ -570,7 +571,7 @@ def main():
'train_res': args['train_res'],
}
trainer = Trainer(training_parameters, name=f'stylegan2-test-10-{args["train_res"]}x{args["train_res"]}')
trainer = Trainer(training_parameters, name=f'stylegan2-test-13-{args["train_res"]}x{args["train_res"]}')
trainer.train(dist_dataset, strategy)
return