gnn with weights

This commit is contained in:
2023-02-25 12:33:18 +00:00
parent 5062ed95d8
commit d8337bfcfe
4 changed files with 17 additions and 13 deletions
+9 -6
View File
@@ -281,7 +281,7 @@ def plot_everything(batch: dict, type: str):
mesh_pred = uv2mesh(pred)
mesh_pred = smoother(mesh_pred, training=False).numpy()
mesh_weigths = uv2mesh_label(weights)
mesh_weigths = smoother_weights(mesh_weigths, training=False).numpy()
mesh_weigths = smoother_weights(mesh_weigths, training=False)
orig_img = uv_orig
uv_static_img = batch['uv_static']
@@ -289,6 +289,8 @@ def plot_everything(batch: dict, type: str):
pred_img = pred
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
# min max normalization of weights
weights = (weights - tf.reduce_min(weights, axis=[1,2,3], keepdims=True)) / (tf.reduce_max(weights, axis=[1,2,3], keepdims=True) - tf.reduce_min(weights, axis=[1,2,3], keepdims=True))
weights_imgs = [wandb.Image(np.concatenate((np.concatenate([weights[b,:,:,i].numpy() for i in range(12)], axis=-1), np.concatenate([weights[b,:,:,i].numpy() for i in range(12,24)], axis=-1)), axis=0)) for b in range(4)]
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
@@ -555,8 +557,8 @@ if __name__ == '__main__':
is_cape=config['is_cape'])
model = CVAE(config)
smoother = Smoother(features=16, activation_middle='tanh', activation_final='tanh', learn_weights=False)
smoother_weights = Smoother(features=16, activation_middle=None, activation_final='softmax', learn_weights=True)
smoother = Smoother(features=24, activation_middle='tanh', activation_final='tanh', learn_weights=False)
smoother_weights = Smoother(features=3, activation_middle=None, activation_final='softmax', learn_weights=True)
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
balancing_activated_step = 0
@@ -581,9 +583,10 @@ if __name__ == '__main__':
train_loss = train_step(batch)
data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
valid_loss = valid_step(data_valid)
if step % 25 == 0:
data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
valid_loss = valid_step(data_valid)
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):