view training trials

models/view_encoder.py

@@ -22,43 +22,73 @@ class ViewEncoder(nn.Module):
         super(ViewEncoder, self).__init__()
         self.bottleneck_size = bottleneck_size
         
-        # Encoder
+        self.encoder = nn.ModuleList([
-        self.encoder = nn.Sequential(
+            nn.Sequential(
-            nn.Linear(3*3, 64),
+                nn.Linear(3*3, 64),
-            nn.ReLU(),
+                nn.BatchNorm1d(64),
-            nn.Linear(64, 128),
+                nn.ReLU(),),
-            nn.ReLU(),
+            nn.Sequential(
-            nn.Linear(128, 256),
+                nn.Linear(64, 128),
-            nn.ReLU(),
+                nn.BatchNorm1d(128),
-            nn.Linear(256, 512),
+                nn.ReLU(),),
-            nn.ReLU(),
+            nn.Sequential(
-            nn.Linear(512, bottleneck_size),
+                nn.Linear(128, 256),
-            nn.ReLU()
+                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(
+        self.decoder = nn.ModuleList([
-            nn.Linear(bottleneck_size, 512),
+            nn.Sequential(
-            nn.ReLU(),
+                nn.Linear(bottleneck_size, 512),
-            nn.Linear(512, 256),
+                nn.BatchNorm1d(512),
-            nn.ReLU(),
+                nn.ReLU(),),
-            nn.Linear(256, 128),
+            nn.Sequential(
-            nn.ReLU(),
+                nn.Linear(512*2, 256),
-            nn.Linear(128, 64),
+                nn.BatchNorm1d(256),
-            nn.ReLU(),
+                nn.ReLU(),),
-            nn.Linear(64, 3*3),
+            nn.Sequential(
-            nn.Tanh(),
+                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__':
     # test the model

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
+        # Compute the Gram-Schmidt process to orthogonalize the columns of the matrix
-        matrix = tensor.view(-1, tensor.size(-1)//2, tensor.size(-1)//2)
+        matrix_size = np.sqrt(tensor.size(-1)).astype(np.int64)
-        # matrix = matrix / torch.norm(matrix, dim=1, keepdim=True)  # Normalize the rows
+        matrix = tensor.view(-1, matrix_size, matrix_size)
+        q, _ = torch.linalg.qr(matrix)
 
-        # Compute the Gram matrix
+        # Compute the Frobenius norm of the difference between the original matrix and the orthogonalized matrix
-        gram_matrix = torch.matmul(matrix, matrix.t())
+        loss = torch.norm(matrix - q, p='fro')
-
-        # 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
     
+# 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()