From f74cb4e92b5d8c9bdf63f07308b058df8f34405a Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 17 Jun 2023 16:11:55 +0000 Subject: [PATCH] feature manipulator pipeline --- scripts/pose_manipulator_train.py | 199 ++++++++++++++++++++++++++++++ 1 file changed, 199 insertions(+) create mode 100644 scripts/pose_manipulator_train.py diff --git a/scripts/pose_manipulator_train.py b/scripts/pose_manipulator_train.py new file mode 100644 index 0000000..e57261f --- /dev/null +++ b/scripts/pose_manipulator_train.py @@ -0,0 +1,199 @@ +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 torch.profiler import profile, record_function, ProfilerActivity +# from torchgeometry.losses.dice import DiceLoss +# from torchmetrics import Dice + +from data.pose_dataset import PoseDataset +from utils.utils import Accumulator +from utils.checkpoints import CheckpointSaver, load_model +from models.pose_manipulator import PoseManipulator +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) + + +name = 'pose_manipulator_1' + +def init_wandb(): + run = wandb.init(project='reconstruction', name=name, notes='now with random poses, 10k ae training', dir='/data-local/127/temp', tags=['pose']) + artifact = wandb.Artifact('code', type='code') + artifact.add_file('scripts/pose_manipulator_train.py') + artifact.add_file('models/pose_manipulator.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) + + +def train_step(batch): + model.train() + + pred_pose, pred_mask, expected_features, modified_features = model([batch['mask'], batch['pose'], batch['mask_out'], batch['pose_out'], batch['angle']]) + + mse_features = loss_fn_features(expected_features, modified_features) + loss: torch.Tensor = mse_features + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + with torch.inference_mode(): + mse = loss_fn(pred_pose, batch['pose_out']) + mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1)) + + return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask, 'train/mse_features': mse_features } + + +def valid_step(batch): + model.eval() + with torch.inference_mode(): + pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']]) + + mse = loss_fn(pred_pose, batch['pose_out']) + mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1)) + symm_rotation = loss_fn_rotation(pred_rotation) + loss = mse + + return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation } + + +def plot_everything(batch, type: str): + model.eval() + with torch.inference_mode(): + pred_pose, pred_mask, _, _ = model([batch['mask'], batch['pose'], batch['mask_out'], batch['pose_out'], batch['angle']]) + + 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)) + pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)) + + masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy() + poses = torch.concat((pose, pose_out, pred_pose), axis=2).cpu().numpy() + masks_imgs = [wandb.Image((masks[i] + 1) / 2.0) for i in range(pred_pose.shape[0])] + poses_imgs = [wandb.Image((poses[i] + 1) / 2.0) for i in range(pred_pose.shape[0])] + + return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, } } + + +def train(): + 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 < 500 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)) } } + + train_accumulator.reset() + + # figure out when to log when using validation set as well + if not args['debug']: + wandb.log(to_log, step=iteration) + # print(f'Iteration {iteration} finished') + + if iteration % 10000 == 0: + ckpt_saver(model.manipulator, iteration, optimizer, scaler=None) + + iteration += 1 + + 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) + + +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 = 4 + use_cache = False + + train_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=False) + valid_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=False) + 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: PoseManipulator = PoseManipulator(device=device, ae_checpoint='checkpoints/reconstruction/pose_ae_2').to(device) + + loss_fn = nn.MSELoss() + loss_fn_features = nn.MSELoss() + loss_fn_mask = nn.BCELoss(reduction='mean') + 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) + + ckpt_saver = CheckpointSaver('checkpoints/reconstruction/' + name, top_n=5) + + # Use the profiler to track memory usage + # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: + # with record_function("model_inference"): + + if not args['debug']: + # model = torch.compile(model) + wandb.watch(model, log='all', log_freq=1000, log_graph=True) + + train() + + # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10)) + + pass \ No newline at end of file