From 2d78839f69647cbd45900bd51034c4a2ba7f92ae Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 6 Jun 2023 22:18:50 +0200 Subject: [PATCH] view model training --- data/view_dataset.py | 55 ++++++++++ models/view_encoder.py | 78 ++++++++++++++ scripts/view_encoder_train.py | 187 ++++++++++++++++++++++++++++++++++ 3 files changed, 320 insertions(+) create mode 100644 data/view_dataset.py create mode 100644 models/view_encoder.py create mode 100644 scripts/view_encoder_train.py diff --git a/data/view_dataset.py b/data/view_dataset.py new file mode 100644 index 0000000..7685cfc --- /dev/null +++ b/data/view_dataset.py @@ -0,0 +1,55 @@ +import os +import sys +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" +# os.environ["CUDA_VISIBLE_DEVICES"] = '3' + +import torch +from torch.utils.data import Dataset +import numpy as np +from tqdm import tqdm +from time import time +from cloth3d.read import DataReader +from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image +from superSMPL.smpl_torch import SMPL +from data.pose_renderer import PoseRenderer + + +class Dataset(Dataset): + def __init__(self, device='cuda'): + self.device = device + self.smpl = SMPL(path='superSMPL', device=device) + + + def __len__(self): + return int(1e6) + + + def __getitem__(self, idx): + # generate a random tenssor of size 3 with values betweeen -pi and pi + angles = torch.rand(3) * 2 * np.pi - np.pi + angles = angles.to(self.device) + matrix = self.smpl.rodrigues(angles[None, None, :]).to(self.device)[0] + + data = { + 'angles': angles, + 'matrix': matrix, + } + + + return data + + +if __name__ == '__main__': + from torch.utils.data import DataLoader + + device = "cuda" if torch.cuda.is_available() else "cpu" + dataset = Dataset(device=device) + train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) + + for batch in train_dataloader: + batch + pass + + pass \ No newline at end of file diff --git a/models/view_encoder.py b/models/view_encoder.py new file mode 100644 index 0000000..7e264bf --- /dev/null +++ b/models/view_encoder.py @@ -0,0 +1,78 @@ +import os +import sys + +os.environ["CUDA_VISIBLE_DEVICES"] = '0' +# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + +import torch +import numpy as np +from torch import nn, Tensor +from torch.nn import functional as F +from functools import partial +from typing import Any, Callable, List, Optional, Sequence + +torch.set_float32_matmul_precision('high') + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + + + +class ViewEncoder(nn.Module): + def __init__(self, bottleneck_size): + super(ViewEncoder, self).__init__() + self.bottleneck_size = bottleneck_size + + # Encoder + self.encoder = nn.Sequential( + nn.Linear(3*3, 64), + nn.ReLU(), + nn.Linear(64, 128), + nn.ReLU(), + nn.Linear(128, 256), + nn.ReLU(), + nn.Linear(256, 512), + nn.ReLU(), + nn.Linear(512, bottleneck_size), + nn.ReLU() + ) + + # Decoder + self.decoder = nn.Sequential( + nn.Linear(bottleneck_size, 512), + nn.ReLU(), + nn.Linear(512, 256), + nn.ReLU(), + nn.Linear(256, 128), + nn.ReLU(), + nn.Linear(128, 64), + nn.ReLU(), + nn.Linear(64, 3*3), + nn.Tanh(), + ) + + def forward(self, x): + # Encode + matrix = self.encoder(x.view(-1, 3*3)) + + # Decode + x = self.decoder(matrix) + x = x * np.pi + return x.view(-1, 3, 3), matrix + + +if __name__ == '__main__': + # test the model + model = ViewEncoder(bottleneck_size=128).to('cuda') + # model = torch.compile(model) + print(model) + print(sum(p.numel() for p in model.parameters() if p.requires_grad)) + print(sum(p.numel() for p in model.parameters())) + + x = torch.rand((4, 3, 3)).to('cuda') + output = model(x) + + # from torchviz import make_dot + # graph = make_dot(output, params=dict(model.named_parameters())) + # graph.render("model_graph", format="pdf") + + pass \ No newline at end of file diff --git a/scripts/view_encoder_train.py b/scripts/view_encoder_train.py new file mode 100644 index 0000000..9b5033e --- /dev/null +++ b/scripts/view_encoder_train.py @@ -0,0 +1,187 @@ +import os +import sys + +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = '0' #config['gpu'] +import torch + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') + +import importlib +import argparse +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 + + +import wandb +import numpy as np +import torch.optim as optim +import torch.nn as nn +from torch.utils.data import DataLoader +# from torchgeometry.losses.dice import DiceLoss +# from torchmetrics import Dice + +from data.view_dataset import Dataset +from utils.utils import Accumulator +from models.view_encoder import ViewEncoder +import plotly.express as ex + + +def init_wandb(): + run = wandb.init(project='reconstruction', name='view_1', notes='', dir='/tmp', tags=['view']) + artifact = wandb.Artifact('code', type='code') + artifact.add_file('scripts/view_encoder_train.py') + artifact.add_file('models/view_encoder.py') + artifact.add_file('data/view_dataset.py') + config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] + for config_file_name in config_file_names: + file_name = '/'.join(config_file_name.split('.')) + '.py' + if os.path.exists(file_name): + artifact.add_file(file_name) + run.log_artifact(artifact) + + +class OrthogonalLoss(nn.Module): + def __init__(self): + super(OrthogonalLoss, self).__init__() + + def forward(self, tensor): + # Reshape the input tensor into a square matrix + matrix = tensor.view(-1, tensor.size(-1)//2, tensor.size(-1)//2) + # matrix = matrix / torch.norm(matrix, dim=1, keepdim=True) # Normalize the rows + + # Compute the Gram matrix + gram_matrix = torch.matmul(matrix, matrix.t()) + + # Compute the Frobenius norm of the difference between the Gram matrix and the identity matrix + loss = torch.norm(gram_matrix - torch.eye(gram_matrix.size(0)).to(tensor.device)) + + return loss + + +def train_step(batch): + model.train() + + pred, matrix = model(batch['matrix']) + + mse = loss_fn(pred, batch['matrix']) + ortho_loss = ortho_loss_fn(matrix) + loss = mse + ortho_loss + + # TODO: calculate max of ortho_loss + + # print(f'Loss: {loss} MSE: {mse}') + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + + return { 'train/loss': loss, 'train/mse': mse, 'train/ortho_loss': ortho_loss } + + +def valid_step(batch): + model.eval() + with torch.inference_mode(): + pred, matrix = model(batch['matrix']) + + mse = loss_fn(pred, batch['matrix']) + ortho_loss = ortho_loss_fn(matrix) + loss = mse + ortho_loss + + + return { 'valid/loss': loss, 'valid/mse': mse } + + +def plot_everything(batch, type): + model.eval() + with torch.inference_mode(): + pred = model(batch['matrix']) + + matrix = batch['matrix'].cpu().numpy() + pred = pred.detach().cpu().numpy() + + matrix_str = [np.array2string(matrix[i], precision=2, suppress_small=True) for i in range(4)] + pred_str = [np.array2string(pred[i], precision=2, suppress_small=True) for i in range(4)] + + logs = { f'{type}/matrix': wandb.Table(data=[[matrix_str, pred_str]], columns=['matrix', 'pred']), } + + return logs + + +if __name__ == '__main__': + # wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d') + if not args['debug']: + init_wandb() + + device = "cuda" if torch.cuda.is_available() else "cpu" + + batch_size = 128 + use_cache = False + + train_dataset = Dataset(device=device) + valid_dataset = Dataset(device=device,) + train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0) + valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() + + model = ViewEncoder(bottleneck_size=1024).to(device) + + loss_fn = nn.MSELoss() + ortho_loss_fn = OrthogonalLoss() + optimizer = optim.Adam(model.parameters(), lr=0.001) + # lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False) + # optimizer = optim.SGD(model.parameters(), lr=0.1) + + iteration = 0 + epoch = 0 + while True: + + for batch in train_dataloader: + train_loss = train_step(batch) + train_accumulator.update(train_loss) + + if iteration % 1000 == 0 or (iteration < 1000 and iteration % 100 == 0): + to_log = dict({ **train_loss, **{ 'current_step': iteration } }) + if iteration % 1000 == 0 or (iteration < 1000 and iteration % 1000 == 0): + train_figs = plot_everything(batch, 'train') + to_log = { **to_log, **train_figs } + # if iteration == 0: + # to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } + + if not args['debug']: + wandb.log(to_log, step=iteration) + train_accumulator.reset() + iteration += 1 + # print(f'Iteration {iteration} finished') + + epoch += 1 + print(f'Epoch {epoch} finished') + + # if not iteration % 1000 == 0: + # continue + + # for batch in valid_dataloader: + # valid_loss = valid_step(batch) + # valid_accumulator.update(valid_loss) + + + # to_log = dict({ **valid_loss, **{ 'current_step': iteration } }) + # valid_figs = plot_everything(batch, 'train') + # to_log = { **to_log, **valid_figs } + # valid_accumulator.reset() + + # if not args['debug']: + # wandb.log(to_log, step=iteration) + + + + + pass \ No newline at end of file