From b4be8bf80dfea34d34fd4eceb28c6be24f28e755 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 6 Apr 2023 17:25:57 +0000 Subject: [PATCH] working static ae training --- configs/ae_static.py | 5 +-- configs/ae_static_orig.py | 71 +++++++++++++++++++++++++++++++++++++ configs/ae_static_orig_2.py | 71 +++++++++++++++++++++++++++++++++++++ models/base.py | 3 +- scripts/ae_static_train.py | 47 ++++++++++++------------ 5 files changed, 172 insertions(+), 25 deletions(-) create mode 100644 configs/ae_static_orig.py create mode 100644 configs/ae_static_orig_2.py diff --git a/configs/ae_static.py b/configs/ae_static.py index d5a9873..ec784a6 100644 --- a/configs/ae_static.py +++ b/configs/ae_static.py @@ -2,11 +2,13 @@ from utils.convnext_utils import CNBlockConfig config = dict() -config['gpu'] = '5' +config['gpu'] = '1' config['name'] = 'static_torch_1' config['notes'] = ' ' config['tags'] = ['torch', 'static'] +config['pretrained_model'] = 'checkpoints/iccv/static_torch_1' + config['batch_size'] = 4 config['train_dir'] = '/data-local/127/train' config['valid_dir'] = '/data-local/127/valid' @@ -60,7 +62,6 @@ config['cond_ae']['activation'] = 'tanh' config['lr'] = 0.0005 config['cond_ae']['upscale_output'] = False -config['pretrained_model'] = '' config['signature'] = { 'id', 'frame_nr', diff --git a/configs/ae_static_orig.py b/configs/ae_static_orig.py new file mode 100644 index 0000000..5a6080a --- /dev/null +++ b/configs/ae_static_orig.py @@ -0,0 +1,71 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '2' +config['name'] = 'static_torch_orig_1' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/train' +config['valid_dir'] = '/data-local/127/valid' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = False +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'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 27), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 27), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 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.0005 +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/configs/ae_static_orig_2.py b/configs/ae_static_orig_2.py new file mode 100644 index 0000000..c02c91a --- /dev/null +++ b/configs/ae_static_orig_2.py @@ -0,0 +1,71 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '3' +config['name'] = 'static_torch_orig_2' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] + +config['batch_size'] = 4 +config['train_dir'] = '/data-local/127/train' +config['valid_dir'] = '/data-local/127/valid' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = None +config['valid_frames'] = None + +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = True +config['is_cape'] = False +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'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 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.0005 +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/models/base.py b/models/base.py index f918b3e..cd9682b 100644 --- a/models/base.py +++ b/models/base.py @@ -69,7 +69,8 @@ class Decoder(nn.Module): self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_classes, device=device, reversed=True) - self.softmax = nn.Softmax(dim=1) + if self.is_label: + self.softmax = nn.Softmax(dim=1) def forward(self, inputs): if len(inputs) == 2: diff --git a/scripts/ae_static_train.py b/scripts/ae_static_train.py index 705a2ba..6986bcc 100644 --- a/scripts/ae_static_train.py +++ b/scripts/ae_static_train.py @@ -13,9 +13,8 @@ import wandb import random import numpy as np import torch -from torch import nn -from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM -from custom_utils.custom_utils import UV2Mesh, get_summary_html, load_model, maskFacesBatch, get_summary_str +# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM +from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, get_summary_str import plotly.graph_objects as go import plotly.express as ex import torch.optim as optim @@ -26,15 +25,18 @@ else: from losses.losses import edge_loss from data.dataset import Dataset -from torch.utils.data import DataLoader -from custom_utils.custom_utils import batch_to_dict, plot_3d +# from torch.utils.data import DataLoader +from custom_utils.custom_utils import plot_3d from models.cond_encoder import AutoEncoder -from utils.checkpoints import CheckpointSaver +from utils.checkpoints import CheckpointSaver, load_model from custom_utils.custom_utils import Accumulator +from utils.utils import InfiniteDataLoader as DataLoader + +from pytorch_msssim import MS_SSIM def init_wandb(): - run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp/wandb', tags=config['tags']) + run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/ae_static_train.py') artifact.add_file('losses/losses.py') @@ -57,19 +59,21 @@ def train_step(batch): model.train() pred = model([batch['uv_static'], batch['uv_mask']]) - mesh_pred = uv2mesh(pred) - mesh_orig = batch['mesh_static'] - - e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1)) ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1)) loss = 1e3 * (1 - ms_ssim) + # loss = l1_loss optimizer.zero_grad() loss.backward() optimizer.step() - metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) + model.eval() + with torch.inference_mode(): + mesh_pred = uv2mesh(pred) + mesh_orig = batch['mesh_static'] + l1_loss = torch.mean(torch.sum((torch.abs(batch['uv_static'] - pred)), axis=1)) + e_loss = 1e-2 * torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } @@ -151,8 +155,6 @@ def train(): ckpt_saver_enc(model.encoder, step, optimizer) ckpt_saver_dec(model.decoder, step, optimizer) - print(step) - step += 1 @@ -170,24 +172,25 @@ if __name__ == '__main__': skirt_modifier = '_skirt' if config['is_skirt'] else '' # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt') - dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames']) - dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames']) - train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=False) - valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=False) + 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']) + 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() valid_accumulator = Accumulator() model : AutoEncoder = AutoEncoder(config['cond_ae']) optimizer = optim.Adam(model.parameters(), lr=config['lr']) - _MSSSIM = MSSSIM(data_range=2).to(config['device']) + # _MSSSIM = MSSSIM(data_range=2).to(config['device']) + _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3) if config['pretrained_model'] != '': load_model(model.encoder, config['pretrained_model'] + '_enc') load_model(model.decoder, config['pretrained_model'] + '_dec') - ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=20) - ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=20) + ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=5) + ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=5) run_name = config['name'] print(f'started training: {run_name}')