218 lines
7.2 KiB
Python
218 lines
7.2 KiB
Python
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 |