From d9021672e4359201b7087352a940bb52da8f8f11 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 17 Jun 2023 16:12:15 +0000 Subject: [PATCH] feature manipulator pipeline --- models/pose_manipulator.py | 79 ++++++++++++ scripts/pose_manipulator_train.py | 199 ++++++++++++++++++++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 models/pose_manipulator.py create mode 100644 scripts/pose_manipulator_train.py diff --git a/models/pose_manipulator.py b/models/pose_manipulator.py new file mode 100644 index 0000000..38d49e4 --- /dev/null +++ b/models/pose_manipulator.py @@ -0,0 +1,79 @@ +import os +import sys + +# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" +# os.environ["CUDA_VISIBLE_DEVICES"] = '0' + +import torch +from torch import nn, Tensor +from torch.nn import functional as F +from functools import partial +from typing import Any, Callable, List, Optional, Sequence + +torch.set_float32_matmul_precision('high') + +sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') +from models.convnext import CNBlock, LayerNorm2d +from models.pose_ae import PoseAE +from utils.checkpoints import load_model + + +class FeatureManipulator(nn.Module): + def __init__( + self, + features: int = 256, + device: str = 'cuda', + ) -> None: + super().__init__() + + norm_layer = partial(LayerNorm2d, eps=1e-6) + self.last = nn.Sequential( + norm_layer(features + 1), + nn.Conv2d(features + 1, features, kernel_size=3, padding='same'), + ).to(device) + + stage: List[nn.Module] = [] + for _ in range(10): + stage.append(CNBlock(features + 1, 1.0, 0)) + self.stage = nn.Sequential(*stage).to(device) + + def _forward_impl(self, x) -> Tensor: + x = self.stage(x) + x = self.last(x) + return x + + def forward(self, x) -> Tensor: + return self._forward_impl(x) + + +class PoseManipulator(nn.Module): + def __init__(self, features=256, device='cuda', ae_checpoint=None): + super().__init__() + + self.pose_ae: PoseAE = PoseAE(features=features, device=device) + iteration = load_model(self.pose_ae.encoder, ae_checpoint + '_enc') + load_model(self.pose_ae.decoder, ae_checpoint + '_dec') + for param in self.pose_ae.encoder.parameters(): + param.requires_grad = False + for param in self.pose_ae.decoder.parameters(): + param.requires_grad = False + + print(f'Loaded pose_ae checkpoint at iteration: {iteration}') + + self.manipulator = FeatureManipulator(features=features, device=device) + + + def forward(self, input) -> torch.Tensor: + mask, pose_in, mask_out, pose_out, angle = input + original_features, _ = self.pose_ae.encoder(torch.cat([mask, pose_in], dim=1)) + expected_features, _ = self.pose_ae.encoder(torch.cat([mask_out, pose_out], dim=1)) + + tiled_angle = torch.ones_like(original_features[:, :1, :, :]) * angle[:,:,None,None] + modified_features = self.manipulator(torch.cat((original_features, tiled_angle), dim=1)) + pose_pred, mask_pred = self.pose_ae.decoder(modified_features) + + return pose_pred, mask_pred, expected_features, modified_features + + +if __name__ == '__main__': + pass \ No newline at end of file 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