diff --git a/models/view_encoder.py b/models/view_encoder.py index 7e264bf..a41d964 100644 --- a/models/view_encoder.py +++ b/models/view_encoder.py @@ -22,42 +22,72 @@ class ViewEncoder(nn.Module): 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() - ) + self.encoder = nn.ModuleList([ + nn.Sequential( + nn.Linear(3*3, 64), + nn.BatchNorm1d(64), + nn.ReLU(),), + nn.Sequential( + nn.Linear(64, 128), + nn.BatchNorm1d(128), + nn.ReLU(),), + nn.Sequential( + nn.Linear(128, 256), + nn.BatchNorm1d(256), + nn.ReLU(),), + nn.Sequential( + nn.Linear(256, 512), + nn.BatchNorm1d(512), + nn.ReLU(),), + nn.Sequential( + nn.Linear(512, bottleneck_size), + nn.BatchNorm1d(bottleneck_size),) + ]) # 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(), - ) + self.decoder = nn.ModuleList([ + nn.Sequential( + nn.Linear(bottleneck_size, 512), + nn.BatchNorm1d(512), + nn.ReLU(),), + nn.Sequential( + nn.Linear(512*2, 256), + nn.BatchNorm1d(256), + nn.ReLU(),), + nn.Sequential( + nn.Linear(256*2, 128), + nn.BatchNorm1d(128), + nn.ReLU(),), + nn.Sequential( + nn.Linear(128*2, 64), + nn.BatchNorm1d(64), + nn.ReLU(),), + nn.Sequential( + nn.Linear(64*2, 3*3), + nn.Tanh(),), + ]) def forward(self, x): # Encode - matrix = self.encoder(x.view(-1, 3*3)) + f1 = self.encoder[0](x.view(-1, 3*3)) + f2 = self.encoder[1](f1) + f3 = self.encoder[2](f2) + f4 = self.encoder[3](f3) + matrix = self.encoder[4](f4) + + d1 = self.decoder[0](matrix) + d2 = self.decoder[1](torch.cat((d1, f4), dim=1)) + d3 = self.decoder[2](torch.cat((d2, f3), dim=1)) + d4 = self.decoder[3](torch.cat((d3, f2), dim=1)) + x = self.decoder[4](torch.cat((d4, f1), dim=1)) - # Decode - x = self.decoder(matrix) x = x * np.pi return x.view(-1, 3, 3), matrix + + # # Decode + # x = self.decoder(matrix) + # x = x * np.pi + # return x.view(-1, 3, 3), matrix if __name__ == '__main__': diff --git a/scripts/view_encoder_train.py b/scripts/view_encoder_train.py index 9b5033e..ef9b424 100644 --- a/scripts/view_encoder_train.py +++ b/scripts/view_encoder_train.py @@ -1,14 +1,10 @@ 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 +import importlib + parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) parser.add_argument('--config', type=str) @@ -17,10 +13,17 @@ args = vars(parser.parse_args()) # config_module = importlib.import_module(args['config']) # config = config_module.config +if args['debug']: + pass +os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" +os.environ["CUDA_VISIBLE_DEVICES"] = '0' +if args['debug']: + os.environ["CUDA_LAUNCH_BLOCKING"] = '1' import wandb import numpy as np +import torch import torch.optim as optim import torch.nn as nn from torch.utils.data import DataLoader @@ -34,7 +37,7 @@ import plotly.express as ex def init_wandb(): - run = wandb.init(project='reconstruction', name='view_1', notes='', dir='/tmp', tags=['view']) + run = wandb.init(project='reconstruction', name='view_7', notes='add skip connections. no activation as bottleneck. different ortholoss', dir='/data-local/127/temp', tags=['view']) artifact = wandb.Artifact('code', type='code') artifact.add_file('scripts/view_encoder_train.py') artifact.add_file('models/view_encoder.py') @@ -47,23 +50,40 @@ def init_wandb(): run.log_artifact(artifact) -class OrthogonalLoss(nn.Module): +class OrthogonalLoss(torch.nn.Module): def __init__(self): - super(OrthogonalLoss, self).__init__() + super().__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-Schmidt process to orthogonalize the columns of the matrix + matrix_size = np.sqrt(tensor.size(-1)).astype(np.int64) + matrix = tensor.view(-1, matrix_size, matrix_size) + q, _ = torch.linalg.qr(matrix) - # 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)) + # Compute the Frobenius norm of the difference between the original matrix and the orthogonalized matrix + loss = torch.norm(matrix - q, p='fro') return loss +# class OrthogonalLoss(nn.Module): +# def __init__(self): +# super(OrthogonalLoss, self).__init__() + +# def forward(self, tensor: torch.Tensor) -> torch.Tensor: +# # Reshape the input tensor into a square matrix +# matrix_size = np.sqrt(tensor.size(-1)).astype(np.int64) +# matrix = tensor.view(-1, matrix_size, matrix_size) +# # matrix = matrix / torch.norm(matrix, dim=1, keepdim=True) # Normalize the rows + +# # Compute the Gram matrix +# gram_matrix = torch.bmm(matrix, matrix.transpose(1, 2)) + +# # Compute the Frobenius norm of the difference between the Gram matrix and the identity matrix +# identity_matrix = torch.eye(matrix_size, device=tensor.device).unsqueeze(0).expand(batch_size, -1, -1) +# loss = torch.norm(gram_matrix - identity_matrix, p='fro') + +# return loss + def train_step(batch): model.train() @@ -72,7 +92,7 @@ def train_step(batch): mse = loss_fn(pred, batch['matrix']) ortho_loss = ortho_loss_fn(matrix) - loss = mse + ortho_loss + loss: torch.Tensor = mse + ortho_loss # TODO: calculate max of ortho_loss @@ -83,6 +103,7 @@ def train_step(batch): optimizer.step() + return { 'train/loss': loss, 'train/mse': mse, 'train/ortho_loss': ortho_loss } @@ -93,31 +114,41 @@ def valid_step(batch): mse = loss_fn(pred, batch['matrix']) ortho_loss = ortho_loss_fn(matrix) - loss = mse + ortho_loss + loss = mse #+ ortho_loss - return { 'valid/loss': loss, 'valid/mse': mse } + return { 'valid/loss': loss, 'valid/mse': mse, 'valid/ortho_loss': ortho_loss } def plot_everything(batch, type): model.eval() with torch.inference_mode(): - pred = model(batch['matrix']) + pred, pred_matrix = model(batch['matrix']) matrix = batch['matrix'].cpu().numpy() pred = pred.detach().cpu().numpy() - + matrix_size = np.sqrt(pred_matrix.size(-1)).astype(np.int64) + pred_matrix = pred_matrix.view(-1, matrix_size, matrix_size) + pred_matrix = pred_matrix.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']), } + # plot pred_matrix as image where each matrix element is a 8x8 pixel square + + + + + matrix_img = [wandb.Image((np.resize(matrix[i], (matrix_size*8,matrix_size*8)))) for i in range(4)] + + logs = { f'{type}/matrix': wandb.Table(data=[[matrix_str, pred_str]], columns=['matrix', 'pred']), f'{type}/matrix_img': matrix_img } return logs if __name__ == '__main__': - # wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d') if not args['debug']: + wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d') init_wandb() device = "cuda" if torch.cuda.is_available() else "cpu" @@ -132,7 +163,7 @@ if __name__ == '__main__': train_accumulator = Accumulator() valid_accumulator = Accumulator() - model = ViewEncoder(bottleneck_size=1024).to(device) + model = ViewEncoder(bottleneck_size=32**2).to(device) loss_fn = nn.MSELoss() ortho_loss_fn = OrthogonalLoss()