From 5f34b667fcaa8244d7d1ca5be0e5bfde5dde6f2d Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Mon, 12 Jun 2023 13:59:53 +0000 Subject: [PATCH] pose pipeline trial --- data/pose_dataset.py | 11 +++ models/pose_encoder.py | 15 ++-- scripts/pose_encoder_train.py | 129 ++++++++++++++++++++-------------- 3 files changed, 95 insertions(+), 60 deletions(-) diff --git a/data/pose_dataset.py b/data/pose_dataset.py index c67fb0f..41a650d 100644 --- a/data/pose_dataset.py +++ b/data/pose_dataset.py @@ -14,6 +14,7 @@ from cloth3d.read import DataReader from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image from superSMPL.smpl_torch import SMPL from data.pose_renderer import PoseRenderer +import trimesh class Dataset(Dataset): @@ -76,13 +77,20 @@ class Dataset(Dataset): body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False) body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False) + mesh = trimesh.Trimesh(vertices=body[0][0,:,:].cpu().numpy(), faces=self.smpl.f_all, process=False, validate=False) + mesh.apply_transform(trimesh.transformations.rotation_matrix(angle=3.14159, direction=[0, 1, 0], point=[0, 0, 0])) + body_back = torch.tensor(mesh.vertices, device=body[0].device).unsqueeze(0).to(torch.float32) + body_back[0,:,0] = -body_back[0,:,0] + # print(f'smpl took {time() - start} seconds') # start = time() render = self._render_body(body[0])[0,:,:,:3] + render_back = self._render_body(body_back)[0,:,:,:3] render_out = self._render_body(body_out[0])[0,:,:,:3] render = render.permute(2,0,1) + render_back = render_back.permute(2,0,1) render_out = render_out.permute(2,0,1) # import plotly.express as ex @@ -90,9 +98,11 @@ class Dataset(Dataset): # ex.imshow(render_out[0,:,:,:3].cpu().numpy()).show() mask = torch.min(render != 1, dim=0)[0][None,:,:].to(torch.float32) + mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32) mask_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32) render = render * mask + render_back = render_back * mask_back render_out = render_out * mask_out # print(f'render took {time() - start} seconds') @@ -101,6 +111,7 @@ class Dataset(Dataset): 'pose': render, 'shape': shape, 'pose_out': render_out, + 'pose_back': render_back, # 'body': body, # 'body_out': body_out, # 'render': render, diff --git a/models/pose_encoder.py b/models/pose_encoder.py index da2702e..5df0e3c 100644 --- a/models/pose_encoder.py +++ b/models/pose_encoder.py @@ -168,7 +168,7 @@ class PoseEncoder(nn.Module): CNBlockConfig(16, None, 3), ] - self.encoder = ConvNeXt(block_setting, input_features=4, device=device) + self.encoder = ConvNeXt(block_setting, input_features=7, device=device) self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device) self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device) self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device) @@ -177,8 +177,8 @@ class PoseEncoder(nn.Module): def forward(self, input) -> torch.Tensor: - mask, pose_in, pose_out = input - features, _ = self.encoder(torch.cat([mask, pose_in], dim=1)) + mask, pose_in, pose_in_back, pose_out, mask_out = input + features, _ = self.encoder(torch.cat([mask, pose_in, pose_in_back], dim=1)) x = features.view(features.size(0), features.size(1), -1) y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1)) rotation = self.rotation_decoder(y) @@ -188,14 +188,15 @@ class PoseEncoder(nn.Module): # x = torch.einsum('bij,bkj->bik', x, rotation) x = torch.einsum('bij,bjk->bik', x, rotation) x = x.view(x.size(0), x.size(1), 16, 16) - pose, mask_out = self.decoder(x) + pose, mask_out_pred = self.decoder(x) + pose = pose * mask_out # mask_out = torch.round(mask_out) # mask_out = (mask_out > 0.5).to(torch.int64) # pose = pose * (mask_out > 0.5).to(torch.float32) - # return pose, mask_out, rotation - x = pose - return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device) + return pose, mask_out_pred, rotation + # x = pose + # return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device) class DumbPoseEncoder(nn.Module): def __init__(self, in_channels=7, out_channels=3, device='cuda'): diff --git a/scripts/pose_encoder_train.py b/scripts/pose_encoder_train.py index 1023c44..898ef6a 100644 --- a/scripts/pose_encoder_train.py +++ b/scripts/pose_encoder_train.py @@ -26,6 +26,7 @@ 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 @@ -40,7 +41,7 @@ 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']) + run = wandb.init(project='reconstruction', name='pose_2', notes='have the back of the body also as the main input', 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') @@ -69,30 +70,34 @@ class SymmetricLoss(nn.Module): def train_step(batch): model.train() - pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']]) + # pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']]) + 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 = 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}') + # print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}') optimizer.zero_grad() loss.backward() optimizer.step() + with torch.inference_mode(): + mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1)) + symm_rotation = loss_fn_rotation(pred_rotation) + 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']]) + 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, batch['mask_out']) + mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1)) symm_rotation = loss_fn_rotation(pred_rotation) - loss = mse + mse_mask + symm_rotation + loss = mse return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation } @@ -100,24 +105,76 @@ def valid_step(batch): def plot_everything(batch, type: str): model.eval() with torch.inference_mode(): - pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']]) + pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_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_back = torch.permute(batch['pose_back'], (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) - 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, } } + masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy() + poses = torch.concat((pose, pose_back, 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])] + + # do a batchwise min-max normalization of pred_rotation + pred_rotation_img = pred_rotation.cpu().numpy() + pred_rotation_img = (pred_rotation_img - pred_rotation_img.min(axis=(1, 2,), keepdims=True)) / (pred_rotation_img.max(axis=(1, 2), keepdims=True) - pred_rotation_img.min(axis=(1, 2), keepdims=True)) + pred_rotation_imgs = [wandb.Image(pred_rotation_img[i]) for i in range(pred_pose.shape[0])] + + return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, f'{type}/rotation': pred_rotation_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) + 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) 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" @@ -141,45 +198,11 @@ if __name__ == '__main__': 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)) } } - - 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) - - +# Use the profiler to track memory usage + # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: + # with record_function("model_inference"): + train() + # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10)) pass \ No newline at end of file