view model training

2023-06-06 22:18:50 +02:00
parent 745d18a324
commit 2d78839f69
3 changed files with 320 additions and 0 deletions
@@ -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 
@@ -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
@@ -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