187 lines
6.7 KiB
Python
187 lines
6.7 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 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.pose_dataset import Dataset
|
|
from utils.utils import Accumulator
|
|
from models.pose_encoder import PoseEncoder as PoseEncoder
|
|
from models.convnext import ConvNeXt
|
|
import plotly.express as ex
|
|
|
|
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
|
|
ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
|
|
|
|
|
|
def init_wandb():
|
|
run = wandb.init(project='reconstruction', name='pose_1', notes='', dir='/data-local/127/temp', tags=['pose'])
|
|
artifact = wandb.Artifact('code', type='code')
|
|
artifact.add_file('scripts/pose_encoder_train.py')
|
|
artifact.add_file('models/pose_encoder.py')
|
|
artifact.add_file('data/pose_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 SymmetricLoss(nn.Module):
|
|
def __init__(self):
|
|
super(SymmetricLoss, self).__init__()
|
|
|
|
def forward(self, x):
|
|
# Check the difference between the matrix and its transpose
|
|
diff = x - x.transpose(1, 2)
|
|
# Compute the Frobenius norm of the difference
|
|
loss = torch.norm(diff, p='fro', dim=(1, 2))
|
|
# Return the average loss over the batch
|
|
return torch.mean(loss)
|
|
|
|
|
|
def train_step(batch):
|
|
model.train()
|
|
|
|
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
|
|
|
|
mse = loss_fn(pred_pose, batch['pose_out'])
|
|
mse_mask = 0 #loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
|
|
symm_rotation = 0 #loss_fn_rotation(pred_rotation)
|
|
loss = mse #+ mse_mask + symm_rotation
|
|
print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}')
|
|
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask, 'train/symm_rotation': symm_rotation }
|
|
|
|
|
|
def valid_step(batch):
|
|
model.eval()
|
|
with torch.inference_mode():
|
|
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
|
|
|
|
mse = loss_fn(pred_pose, batch['pose_out'])
|
|
mse_mask = loss_fn_mask(pred_mask, batch['mask_out'])
|
|
symm_rotation = loss_fn_rotation(pred_rotation)
|
|
loss = mse + mse_mask + symm_rotation
|
|
|
|
|
|
return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
|
|
|
|
|
|
def plot_everything(batch, type: str):
|
|
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
|
|
|
|
mask = torch.permute(batch['mask'], (0, 2, 3, 1))
|
|
mask_out = torch.permute(batch['mask_out'], (0, 2, 3, 1))
|
|
pose = torch.permute(batch['pose'], (0, 2, 3, 1))
|
|
pose_out = torch.permute(batch['pose_out'], (0, 2, 3, 1))
|
|
pred_mask = torch.permute(pred_mask, (0, 2, 3, 1)).detach()
|
|
pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)).detach()
|
|
|
|
masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1)
|
|
poses = torch.concat((pose, pose_out, pred_pose), axis=2)
|
|
masks_imgs = [wandb.Image((masks[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
|
|
poses_imgs = [wandb.Image((poses[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
|
|
return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, } }
|
|
|
|
|
|
if __name__ == '__main__':
|
|
if not args['debug']:
|
|
init_wandb()
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
batch_size = 1
|
|
use_cache = False
|
|
|
|
train_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
|
|
valid_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
|
|
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 = PoseEncoder(device=device).to(device)
|
|
|
|
loss_fn = nn.MSELoss()
|
|
loss_fn_mask = nn.BCELoss(reduction='mean')
|
|
loss_fn_rotation = SymmetricLoss()
|
|
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 % 5000 == 0 or (iteration < 5000 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 |