import os import sys sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') import argparse import importlib 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 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 # 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_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') 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(torch.nn.Module): def __init__(self): super().__init__() def forward(self, tensor): # 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 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() pred, matrix = model(batch['matrix']) mse = loss_fn(pred, batch['matrix']) ortho_loss = ortho_loss_fn(matrix) loss: torch.Tensor = 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, 'valid/ortho_loss': ortho_loss } def plot_everything(batch, type): model.eval() with torch.inference_mode(): 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)] # 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__': 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" 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=32**2).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