luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

view training trials

Browse Source
This commit is contained in:
luni
2023-06-07 13:25:42 +00:00
parent dc2bc318a7
commit 9f7444f1e4
2 changed files with 114 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/view_encoder_train.py
+56 -25
View File
@@ -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()