From 70162cf8acfed94fae04df210c7dd140fe7c8c3f Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 21 May 2023 19:45:24 +0200 Subject: [PATCH] adapt ae_body to cape --- configs/ae_body_cape_m.py | 73 +++++++++++++++++++++++++++++++++++++++ losses/losses.py | 26 +++++++++----- scripts/ae_body_train.py | 33 +++++++++--------- 3 files changed, 108 insertions(+), 24 deletions(-) create mode 100644 configs/ae_body_cape_m.py diff --git a/configs/ae_body_cape_m.py b/configs/ae_body_cape_m.py new file mode 100644 index 0000000..b7b0a94 --- /dev/null +++ b/configs/ae_body_cape_m.py @@ -0,0 +1,73 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '4' +config['name'] = 'body_torch_cape_m' +config['notes'] = ' ' +config['tags'] = ['torch', 'body'] +config['script'] = 'scripts/ae_body_train.py' + +config['pretrained_model'] = '' + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/cape_processed/train/male' +config['valid_dir'] = '/data-local/127/cape_processed/test/male' +config['test_dir'] = '/data-local/127/cape_processed/test/male' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = True +config['remove_global_orientation'] = True +config['device'] = 'cuda' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +config['bottleneck_scale'] = 2 + +config['cond_ae'] = dict() +config['cond_ae']['device'] = config['device'] +config['cond_ae']['is_label'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), + # CNBlockConfig(3, 3, 3), +] + +config['cond_ae']['encoder_block_settings'] = block_setting +config['cond_ae']['decoder_block_settings'] = block_setting_reverse +config['cond_ae']['activation'] = 'tanh' + + +config['lr'] = 0.0001 +config['cond_ae']['upscale_output'] = False + +config['signature'] = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_mask', + 'mesh_body_posed', + } \ No newline at end of file diff --git a/losses/losses.py b/losses/losses.py index 2cc7bcb..9fdc46c 100644 --- a/losses/losses.py +++ b/losses/losses.py @@ -24,14 +24,14 @@ F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda') -def _nearest_neighbour(V, B): +def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor): # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function batch_size = V.shape[0] batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1]) batch_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1]) V = V.reshape(-1, V.shape[-1]) B = B.reshape(-1, B.shape[-1]) - idx = knn(V, B, k=1, batch_x=batch_x, batch_y=batch_y)[1] + idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1] idx = (idx - batch_x * 14475).reshape(batch_size, 14475) return idx @@ -64,18 +64,28 @@ edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense()) laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo() edges_skirt: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda') +smpl_to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1) +smpl_to_keep = smpl_to_keep[smpl_to_keep < 6890] +laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep,:][:,smpl_to_keep] +edges_cape = torch.tensor(np.stack(np.nonzero(laplacian_loss_orig)).transpose(), device='cuda') -def get_edge_mask(masks: torch.Tensor): +def get_edge_mask(edges: torch.Tensor, masks: torch.Tensor): if edges.get_device() != masks.get_device(): edges.to(masks.get_device()) return torch.min(masks[:,edges], dim=-1)[0] def edge_loss(x, y, mask): - e_mask = get_edge_mask(mask) - x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:] + if x.shape[1] == 14475: + edges_local = edges + elif x.shape[1] == 3638: + edges_local = edges_cape + else: + raise ValueError("Invalid number of vertices") + e_mask = get_edge_mask(edges_local, mask) + x_e = x[:,edges_local[:,0],:] - x[:,edges_local[:,1],:] x_e = torch.sum(x_e ** 2, -1) x_e = torch.sqrt(x_e + 1e-9) - y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:] + y_e = y[:,edges_local[:,0],:] - y[:,edges_local[:,1],:] y_e = torch.sum(y_e ** 2, -1) y_e = torch.sqrt(y_e + 1e-9) d_e = y_e - x_e @@ -173,8 +183,8 @@ def log_normal_pdf(sample, mean, logvar, raxis=1): def kl_loss(mean, logvar): - kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar) - kl_loss = tf.reduce_mean(kl_loss) + kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar) + kl_loss = torch.mean(kl_loss) kl_loss *= -0.5 return kl_loss diff --git a/scripts/ae_body_train.py b/scripts/ae_body_train.py index 456a999..0cf2bff 100644 --- a/scripts/ae_body_train.py +++ b/scripts/ae_body_train.py @@ -8,14 +8,16 @@ import importlib parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) parser.add_argument('--config', type=str) +parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) config = config_module.config -# from configs.ae_body_1 import config +# from configs.ae_body_cape_m import config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] - +if args['debug']: + os.environ["CUDA_LAUNCH_BLOCKING"] = '1' import time import wandb @@ -155,37 +157,31 @@ def train(): if step == 0: to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - wandb.log(to_log, step=step) + if not args['debug']: + wandb.log(to_log, step=step) train_accumulator.reset() valid_accumulator.reset() print(step) if step % 10000 == 0: - ckpt_saver_enc(model.encoder, step, optimizer) - ckpt_saver_dec(model.decoder, step, optimizer) + ckpt_saver_enc(model.encoder, step, optimizer, scaler=None) + ckpt_saver_dec(model.decoder, step, optimizer, scaler=None) step += 1 if __name__ == '__main__': - # init_wandb() - # device = "cuda" if torch.cuda.is_available() else "cpu" + if not args['debug']: + init_wandb() - # parser = argparse.ArgumentParser(description='') - # parser.add_argument('--name', type=str) - # args = vars(parser.parse_args()) - - resolution_modifier = '' - if config['resolution'] != 512: - resolution_modifier = '_' + str(config['resolution']) skirt_modifier = '_skirt' if config['is_skirt'] else '' uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy') uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2) uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device']) - dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory']) - dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory']) + dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape']) + dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape']) train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True) valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True) train_accumulator = Accumulator() @@ -206,6 +202,11 @@ if __name__ == '__main__': run_name = config['name'] print(f'started training: {run_name}') + + if not args['debug']: + # model = torch.compile(model) + wandb.watch(model, log='all', log_freq=1000, log_graph=True) + train()