This commit is contained in:
2022-06-14 11:06:14 +00:00
parent 245f811d47
commit 2dd059d8e8
+15 -13
View File
@@ -17,13 +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 from losses import collision_loss as collision_penalty
os.environ["CUDA_VISIBLE_DEVICES"] = "1" os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_10(test-13)') run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='stylegan', name='test_conditioned_10(test-13)', notes='same as test-12 but double discriminator feature sizes')
# artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
# artifact.add_file('train.py') artifact.add_file('train.py')
# artifact.add_file('losses.py') artifact.add_file('losses.py')
# artifact.add_file('stylegan2/layers/mapping_block.py') artifact.add_file('stylegan2/layers/mapping_block.py')
# run.log_artifact(artifact) run.log_artifact(artifact)
def initiate_models(g_params, d_params, use_custom_cuda): def initiate_models(g_params, d_params, use_custom_cuda):
discriminator = load_discriminator(d_params, ckpt_dir=None, custom_cuda=use_custom_cuda) discriminator = load_discriminator(d_params, ckpt_dir=None, custom_cuda=use_custom_cuda)
@@ -206,7 +206,7 @@ class Trainer(object):
mesh_pred = uv2mesh(fake_images) mesh_pred = uv2mesh(fake_images)
e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask) e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask)
c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) c_loss = 0.0 #collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask) m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask)
# scale loss # scale loss
@@ -235,7 +235,7 @@ class Trainer(object):
mesh_pred = uv2mesh(fake_images) mesh_pred = uv2mesh(fake_images)
e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask) e_loss = edge_loss(mesh_pred, mesh_static, mesh_mask)
c_loss = collision_penalty(mesh_pred, mesh_body_posed, mesh_mask) c_loss = 0.0 #collision_penalty(mesh_pred, mesh_body_posed, mesh_mask)
m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask) m_loss = 0.0 # uv_mask_loss(fake_images, uv_mask)
# scale loss # scale loss
@@ -515,22 +515,24 @@ def main():
# network params # network params
resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024] resolutions = [4, 8, 16, 32, 64, 128, 256, 512, 1024]
featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16] g_featuremaps = [256, 256, 256, 256, 256, 128, 64, 32, 16]
d_featuremaps = [512, 512, 512, 512, 512, 256, 128, 64, 32]
# 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] # featuremaps = [64, 64, 64, 64, 64, 32, 16, 8, 4]
train_resolutions, train_featuremaps = filter_resolutions_featuremaps(resolutions, featuremaps, args['train_res']) train_resolutions, train_g_featuremaps = filter_resolutions_featuremaps(resolutions, g_featuremaps, args['train_res'])
_, train_d_featuremaps = filter_resolutions_featuremaps(resolutions, d_featuremaps, args['train_res'])
g_params = { g_params = {
'z_dim': 256, 'z_dim': 256,
'w_dim': 256, 'w_dim': 256,
'labels_dim': -1, 'labels_dim': -1,
'n_mapping': 8, 'n_mapping': 8,
'resolutions': train_resolutions, 'resolutions': train_resolutions,
'featuremaps': train_featuremaps, 'featuremaps': train_g_featuremaps,
} }
d_params = { d_params = {
'labels_dim': -1, 'labels_dim': -1,
'resolutions': train_resolutions, 'resolutions': train_resolutions,
'featuremaps': train_featuremaps, 'featuremaps': train_d_featuremaps,
} }
# prepare distribute strategy # prepare distribute strategy