import os os.environ["CUDA_VISIBLE_DEVICES"] = "2" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') if gpus: try: # Currently, memory growth needs to be the same across GPUs for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) logical_gpus = tf.config.experimental.list_logical_devices('GPU') print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs") except RuntimeError as e: # Memory growth must be set before GPUs have been initialized print(e) import time import argparse import wandb import random import numpy as np import scipy as sp from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces import plotly.graph_objects as go import plotly.express as ex from PIL import Image from utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing from load_models import load_generator, load_discriminator from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss # from losses import collision_penalty as collision_penalty from losses import collision_loss as collision_penalty from custom_utils.custom_utils import batch_to_dict, load_model from dataset import DatasetFactory, StaticDataset from custom_utils.smpl import SMPL from models.cvae_pyramid_model import Pyramid name='cvae_pyr_1_d128' def init_wandb(): run = wandb.init(project='vae', name=name, notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \ new implementation for kl loss. no lr schedule. 25 jumpsuits, average reconst_loss before calculation, linearly increasing balancing, up to 1 from 0.1. \ offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \ retrained body encoder with deoriented bodies \ use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\ fixed gender issue and retrained body encoder\ remove shape from unposed garment\ training on 100 garments with offsets over the average of these\ l2 loss on both 128 and full res, bottom predicts base, high res predicts offsets over this') artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_pyramid.py') artifact.add_file('models/cvae_model.py') artifact.add_file('models/cvae_pyramid_model.py') artifact.add_file('dataset.py') artifact.add_file('losses.py') artifact.add_file('convnext/models/convnext_tf.py') run.log_artifact(artifact) uv2mesh = UV2Mesh(res=256) def plot_3d(mesh_gt, mesh_pred_unposed, mesh_pred_unposed_low, faces, collision, collision_faces, collision_color): mesh_pred_unposed[:,2] += 1 mesh_pred_unposed_low[:,2] += 2 collision[:,2] += 3 fig = go.Figure(data=[ go.Mesh3d( x=mesh_gt[:,0], y=mesh_gt[:,2], z=mesh_gt[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2] ), go.Mesh3d( x=mesh_pred_unposed[:,0], y=mesh_pred_unposed[:,2], z=mesh_pred_unposed[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], ), go.Mesh3d( x=mesh_pred_unposed_low[:,0], y=mesh_pred_unposed_low[:,2], z=mesh_pred_unposed_low[:,1], i = faces[:,0], j = faces[:,1], k = faces[:,2], ), go.Mesh3d( x=collision[:,0], y=collision[:,2], z=collision[:,1], 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=3, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,5],), zaxis = dict(nticks=4, range=[-1,1],),) ) return fig def log_normal_pdf(sample, mean, logvar, raxis=1): log2pi = tf.math.log(2. * np.pi) res = tf.reduce_sum(-.5 * ((sample - mean) ** 2. * tf.exp(-logvar) + logvar + log2pi), axis=raxis) return res def kl_loss(mean, logvar): kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) kl_loss = tf.reduce_mean(kl_loss) kl_loss *= -0.5 return kl_loss @tf.function def train_step(batch): with tf.GradientTape() as tape: uv_curvature = tf.one_hot(batch['uv_curvature'], 7) uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low]) mesh_pred_full = uv2mesh(pred_full) mesh_pred_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full) mesh_pred_low = uv2mesh(pred_low) mesh_pred_posed_low = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_low) mesh_orig = batch['mesh_unposed'] shape = batch['pose_shape'][:, 72:] gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) mesh_pred_posed_full = mesh_pred_posed_full + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] mesh_pred_posed_low = mesh_pred_posed_low + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis] e_loss = tf.reduce_mean(edge_loss(mesh_pred_full, mesh_orig, batch['mesh_mask'])) c_loss = tf.reduce_mean(collision_penalty(mesh_pred_posed_full, batch['mesh_body_posed'], batch['mesh_mask'])) l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full)), axis=-1) * batch['uv_mask'], axis=[1, 2]) l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred_full) ** 2), axis=-1) * batch['uv_mask'], axis=[1, 2]) ms_ssim = tf.image.ssim_multiscale(pred_full + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3) kl = kl_loss(mean_full, logvar_full) e_loss_low = tf.reduce_mean(edge_loss(mesh_pred_low, mesh_orig, batch['mesh_mask'])) c_loss_low = tf.reduce_mean(collision_penalty(mesh_pred_posed_low, batch['mesh_body_posed'], batch['mesh_mask'])) l1_loss_low = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low)), axis=-1), axis=[1, 2]) l2_loss_low = 1e3 * tf.reduce_mean(tf.reduce_sum(((tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) - pred_low) ** 2), axis=-1), axis=[1, 2]) ms_ssim_low = tf.image.ssim(pred_low + 1, tf.transpose(batch['uv_unposed' + resolution_modifier], [0,2,3,1]) + 1, max_val=3) kl_low = kl_loss(mean_full, logvar_full) reconstr_loss = l2_loss + l2_loss_low reconstr_loss = tf.reduce_mean(reconstr_loss) loss = reconstr_loss + balancing_coeff * (kl + kl_low) ms_ssim = tf.reduce_mean(ms_ssim) l1_loss = tf.reduce_mean(l1_loss) l2_loss = tf.reduce_mean(l2_loss) mean = tf.reduce_mean(mean_full) logvar = tf.reduce_mean(logvar_full) ms_ssim_low = tf.reduce_mean(ms_ssim_low) l1_loss_low = tf.reduce_mean(l1_loss_low) l2_loss_low = tf.reduce_mean(l2_loss_low) mean_low = tf.reduce_mean(mean_low) logvar_low = tf.reduce_mean(logvar_low) mesh_orig_2 = uv2mesh(batch['uv_unposed']) metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) metric_low = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_low) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred_full) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, 'train/euclid_error_reconst': metric_2, 'train/reconstr_loss': reconstr_loss, 'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar, 'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss, 'train/edge_loss' + resolution_modifier: e_loss_low, 'train/collision_loss' + resolution_modifier: c_loss_low, 'train/l1_loss' + resolution_modifier: l1_loss_low, 'train/euclid_error' + resolution_modifier: metric_low, 'train/ms_ssim' + resolution_modifier: ms_ssim_low, 'train/mean' + resolution_modifier: mean_low, 'train/logvar' + resolution_modifier: logvar_low, 'train/kl_loss' + resolution_modifier: kl_low, 'train/var' + resolution_modifier: np.e ** logvar_low, 'train/l2_loss' + resolution_modifier: l2_loss_low } @tf.function def valid_step(batch): uv_curvature = tf.one_hot(batch['uv_curvature'], 7) pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask'], batch['uv_shaped_body']]) mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) mesh_orig = batch['mesh_unposed'] e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss = 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred)), axis=-1) * batch['uv_mask'], axis=[1, 2]) b_loss = 0 #1e-3 * bend_loss(mesh_pred, smpl_faces, laplacian_loss, batch['mesh_mask']) # ms_ssim = tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2) ms_ssim = tf.image.ssim_multiscale(pred + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=3) fft_loss = 0 #tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(tf.transpose(batch['uv_unposed'], [0,2,3,1]), tf.zeros_like(tf.transpose(batch['uv_unposed'], [0,2,3,1]))))), axis=[1,2,3]) # mesh_orig_2 = uv2mesh(batch['uv_unposed']) metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) metric_2 = 0 #tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) reconstr_loss = l1_loss # logpx_z = -reconstr_loss # logpz = log_normal_pdf(z[:,:model.latent_dim], 0., 0.) # logqz_x = log_normal_pdf(z[:,:model.latent_dim], mean, logvar) # loss = -tf.reduce_mean(logpx_z + balancing_coeff * (logpz - logqz_x)) kl = kl_loss(mean, logvar) # loss = reconstr_loss / (2 * (gamma ** 2)) + kl loss = reconstr_loss + balancing_coeff * kl ms_ssim = tf.reduce_mean(ms_ssim) l1_loss = tf.reduce_mean(l1_loss) fft_loss = tf.reduce_mean(fft_loss) b_loss = tf.reduce_mean(b_loss) reconstr_loss = tf.reduce_mean(reconstr_loss) mean = tf.reduce_mean(mean) logvar = tf.reduce_mean(logvar) return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, 'valid/fft_loss': fft_loss, 'valid/euclid_error_reconst': metric_2, 'train/bend_loss': b_loss, 'train/reconstr_loss': reconstr_loss, 'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar } def plot_everything(batch: dict, type: str): uv_curvature = tf.one_hot(batch['uv_curvature'], 7) uv_curvature_low = tf.one_hot(batch['uv_curvature' + resolution_modifier], 7) pred_full, mean_full, logvar_full, pred_low, mean_low, logvar_low, = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_unposed' + resolution_modifier], batch['uv_static' + resolution_modifier], batch['uv_body_posed' + resolution_modifier], uv_curvature_low]) mesh_pred_full = uv2mesh(pred_full) mesh_pred_low = uv2mesh(pred_low) mesh_orig = batch['mesh_unposed'] orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1]) uv_static_img = tf.transpose(batch['uv_static'], [0,2,3,1]) uv_body_posed_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1]) pred_img = pred_full pred_low_up = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5) img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img, pred_low_up), axis=2) imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] F = np.load("custom_utils/smpl_faces.npy") mask = batch['mesh_mask'].numpy().astype(bool) garment_unposed, faces = maskFacesBatch(mesh_pred_full.numpy(), F, mask) garment_unposed_low, faces = maskFacesBatch(mesh_pred_low.numpy(), F, mask) gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred_full.shape[0])] body = batch['mesh_body_posed'].numpy() mesh_posed_full = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_full).numpy() shape = batch['pose_shape'][:, 72:] gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32) shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])]) garment_posed = [mesh_posed_full[i, mask[i], :] + shape_offsets.numpy()[i, mask[i], :] for i in range(body.shape[0])] collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))] collision_faces = [np.concatenate((F, faces[i] + 14475), axis=0) for i in range(len(garment_posed))] collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))] figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], garment_unposed_low[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) } mean_hist = wandb.Histogram(sequence=mean_full.numpy().mean(axis=0)) logvar_hist = wandb.Histogram(sequence=logvar_full.numpy().mean(axis=0)) return { **figs, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } } def sample(folder): # run = wandb.init(project='vae_plots', name=folder.split('/')[-1]) load_model(model,folder) count = 0 for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) uv_curvature = tf.one_hot(batch['uv_curvature'], 7) nr_of_samples=2 pred = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples) pose, shape = batch['pose_shape'][0, :72], batch['pose_shape'][0, 72:] pose, shape = tf.tile(pose[tf.newaxis,:], [nr_of_samples,1]), tf.tile(shape[tf.newaxis,:], [nr_of_samples,1]) mesh_pred = uv2mesh(tf.transpose(pred, [0,3,1,2])) mesh_posed = smpl.run_garment(pose, shape, mesh_pred).numpy() mesh_orig = batch['mesh_posed'][0].numpy() mesh_orig_reposed = smpl.run_garment(pose[:1], shape[:1], batch['mesh_unposed'][:1]).numpy()[0] body = batch['mesh_body_posed'][0].numpy() body_color = np.ones_like(body, dtype=np.int32)*100 F = np.load("custom_utils/smpl_faces.npy") mask = batch['mesh_mask'].numpy().astype(bool)[0] _, faces = _maskFaces(np.ones((14475, 3)), F, mask) pred_zero = model.sample(batch['uv_static'][0], batch['uv_body_posed'][0], uv_curvature[0], batch['uv_mask'][0], nr_of_samples=nr_of_samples, zero_latent=True) mesh_pred_zero = uv2mesh(tf.transpose(pred_zero, [0,3,1,2])) mesh_posed_zero = smpl.run_garment(pose, shape, mesh_pred_zero).numpy() meshes = [] meshes.append(go.Mesh3d( # 8 vertices of a cube x=mesh_orig_reposed[mask,0], y=mesh_orig_reposed[mask,2], z=mesh_orig_reposed[mask,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] )) for i in range(nr_of_samples): x = mesh_posed[i][mask, :].copy() # x[:, 2] += 1 * (i + 1) x[:, 2] += 1 meshes.append(go.Mesh3d( # 8 vertices of a cube x=x[:,0], y=x[:,2], z=x[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] )) for i in range(nr_of_samples): x = mesh_posed_zero[i][mask, :].copy() # x[:, 2] += 1 * (i + 1) x[:, 2] += 2 meshes.append(go.Mesh3d( # 8 vertices of a cube x=x[:,0], y=x[:,2], z=x[:,1], # i, j and k give the vertices of triangles i = faces[:,0], j = faces[:,1], k = faces[:,2] )) # for i in range(5): # x = body.copy() # x[:, 2] += 1 * i # meshes.append(go.Mesh3d( # # 8 vertices of a cube # x=x[:,0], # y=x[:,2], # z=x[:,1], # # i, j and k give the vertices of triangles # i = F[:,0], # j = F[:,1], # k = F[:,2], # vertexcolor=body_color # )) fig = go.Figure(data=meshes) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=2, z=1)) # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=nr_of_samples-1, z=1)) fig.update_layout( scene = dict( xaxis = dict(nticks=4, range=[-1,1],), # yaxis = dict(nticks=4, range=[-1,nr_of_samples+1],), yaxis = dict(nticks=4, range=[-1,3],), zaxis = dict(nticks=4, range=[-1,1],),) ) # wandb.log({'plot_' + str(count): fig}) fig.show() pred = pred.numpy() frames = [Image.fromarray(((pred[i] + 1) * 128).astype(np.int8), mode='RGB') for i in range(pred.shape[0])] frames[0].save("temp/garment_1.gif", format="GIF", append_images=frames[1:], save_all=True, duration=500, loop=0) if count > 10: break print(count) count += 1 print() if __name__ == '__main__': init_wandb() allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) args = vars(parser.parse_args()) batch_size = 4 low_resolution = 128 resolution_modifier = '_' + str(low_resolution) smpl = SMPL() factory = DatasetFactory() signature_dynamic = { 'id': (), 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), 'uv_body_posed': (3, 512, 256,), 'uv_body_posed' + resolution_modifier: (3, low_resolution, low_resolution//2,), 'mesh_body_posed': (14475, 3,), 'mesh_posed': (14475, 3,), 'mesh_unposed': (14475, 3,), 'uv_unposed': (3, 512, 256,), 'uv_unposed' + resolution_modifier: (3, low_resolution, low_resolution//2,), 'uv_curvature': (512, 256), 'uv_curvature' + resolution_modifier: (low_resolution, low_resolution//2), } signature_static = { 'uv_static': (3, 512, 256,), 'uv_static' + resolution_modifier: (3, low_resolution, low_resolution//2,), # 'mesh_static': (14475, 3,), 'mesh_mask': (14475,), 'uv_mask': (512, 256,), # 'uv_shaped_body': (3, 512, 256,), } dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=batch_size, signature=signature_dynamic, use_filled=True, resolutions=[512, low_resolution]) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=[512, low_resolution]) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=batch_size, signature=signature_dynamic) static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) model = Pyramid(latent_dim=128, low_res=low_resolution) # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) optimizer = tf.optimizers.Adam(learning_rate=0.0001) balancing_coeff = 0.0001 reg_coeff = 1 folder = 'models/ae_logs/' valid_iterator = iter(dataset_valid) gamma = 1 balancing_activated_step = 0 # avg_jumpsuit = np.load('custom_data/jumpsuit_20_avg_uv.npy') # sample('checkpoints/cvae_13') ckpt = tf.train.Checkpoint(model=model) manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) step = 0 for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic)) train_loss = train_step(batch) # data_valid = valid_iterator.get_next() # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) # valid_loss = valid_step(data_valid) valid_loss = {} to_log = { **train_loss, **valid_loss, **{ 'current_step': step} } if step % 2500 == 0: train_figs = plot_everything(batch, 'train') # valid_figs = plot_everything(data_valid, 'valid') valid_figs = {} to_log = { **to_log, **train_figs, **valid_figs } if step % 10000 == 0: manager.save(checkpoint_number=step) if train_loss['train/euclid_error'] < 0.01: if balancing_activated_step == 0: balancing_activated_step = step balancing_coeff = balancing_coeff + 1 / 15000 # if train_loss['train/euclid_error'] < 0.015: # reg_coeff = reg_coeff + 100 / 15000 # if step > 15000 and balancing_coeff <= 1e-3: # balancing_coeff = 1e-2 # if step > 20000 and balancing_coeff <= 1e-3: # balancing_coeff = 1e-2 # if step > 25000 and balancing_coeff <= 1e-2: # balancing_coeff = 1e-1 # if step > 30000 and balancing_coeff <= 1e-1: # balancing_coeff = 1 # gamma = 0.99 * gamma + 0.01 * train_loss['train/reconstr_loss'] to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } } wandb.log(to_log, step = step) print(step) step += 1