From 7a31f0d04c26e6d1e64e8b4ca6f93c212cae68ae Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 16 Jun 2023 22:11:57 +0000 Subject: [PATCH] pose autoencoder pipeline --- data/pose_dataset.py | 61 ++++++++----- models/convnext.py | 3 +- models/pose_ae.py | 183 +++++++++++++++++++++++++++++++++++++++ scripts/pose_ae_train.py | 182 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 405 insertions(+), 24 deletions(-) create mode 100644 models/pose_ae.py create mode 100644 scripts/pose_ae_train.py diff --git a/data/pose_dataset.py b/data/pose_dataset.py index 41a650d..8aeaef9 100644 --- a/data/pose_dataset.py +++ b/data/pose_dataset.py @@ -17,10 +17,11 @@ from data.pose_renderer import PoseRenderer import trimesh -class Dataset(Dataset): - def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True): +class PoseDataset(Dataset): + def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True, ae_mode=False): self.folder = folder self.device = device + self.ae_mode = ae_mode self.frames = read_frames_file(os.path.join(folder, frames_file)) self.reader = DataReader(folder, 'superSMPL') @@ -49,7 +50,7 @@ class Dataset(Dataset): # and rotate it 90 degress along the z axis body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T - out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.get_device())) + out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.device)) return out @@ -68,57 +69,73 @@ class Dataset(Dataset): shape = torch.tensor(self.shapes[frame.id], device=self.device) pose[:3] = 0 pose_out = pose.clone() - pose_out[1] = np.random.uniform(-np.pi, np.pi) + + angle = torch.rand(1) * 2 * np.pi - np.pi + matrix = self.smpl.rodrigues(torch.tensor([0.0, angle.numpy()[0], 0.0], device=self.device)[None, None, :]).to(self.device)[0] + angle = angle.to(self.device) + + pose_out[1] = angle # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi) # generate random number between -np.pi, np.pi # start = time() - body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False) - body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False) + if self.ae_mode: + body = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False) + else: + 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] + # 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_back = self._render_body(body_back)[0,:,:,:3] + if not self.ae_mode: + 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) + # render_back = render_back.permute(2,0,1) + if not self.ae_mode: + render_out = render_out.permute(2,0,1) # import plotly.express as ex # ex.imshow(render[0,:,:,:3].cpu().numpy()).show() # 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) + # mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32) + if not self.ae_mode: + 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 + # render_back = render_back * mask_back + if not self.ae_mode: + render_out = render_out * mask_out # print(f'render took {time() - start} seconds') data = { 'pose': render, 'shape': shape, - 'pose_out': render_out, - 'pose_back': render_back, + # 'pose_back': render_back, # 'body': body, # 'body_out': body_out, # 'render': render, # 'render_out': render_out, 'mask': mask, - 'mask_out': mask_out, + 'angle': angle, + 'matrix': matrix, } + + if not self.ae_mode: + data['pose_out'] = render_out + data['mask_out'] = mask_out if self.cache is not None: self.cache[frame] = data @@ -129,7 +146,7 @@ if __name__ == '__main__': from torch.utils.data import DataLoader device = "cuda" if torch.cuda.is_available() else "cpu" - dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device) + dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device) train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False) for batch in train_dataloader: diff --git a/models/convnext.py b/models/convnext.py index ef5986d..fad998b 100644 --- a/models/convnext.py +++ b/models/convnext.py @@ -159,11 +159,10 @@ class ConvNeXtBase(nn.Module): def _forward_impl(self, x: Tensor) -> Tensor: x = self.stem(x) - features: List[Tensor] = [] x = self.layers(x) x = self.act(x) - return x, features + return x, None def forward(self, x: Tensor) -> Tensor: return self._forward_impl(x) diff --git a/models/pose_ae.py b/models/pose_ae.py new file mode 100644 index 0000000..6fd22ff --- /dev/null +++ b/models/pose_ae.py @@ -0,0 +1,183 @@ +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 ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d + + +class ConvNeXtReversePose(nn.Module): + def __init__( + self, + block_setting: List[CNBlockConfig], + stochastic_depth_prob: float = 0.0, + layer_scale: float = 1e-6, + block: Optional[Callable[..., nn.Module]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + device: str = 'cuda', + use_skip=False, + ) -> None: + super().__init__() + + self.use_skip = use_skip + + if not block_setting: + raise ValueError("The block_setting should not be empty") + elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])): + raise TypeError("The block_setting should be List[CNBlockConfig]") + + if block is None: + block = CNBlock + + if norm_layer is None: + norm_layer = partial(LayerNorm2d, eps=1e-6) + + stage: List[nn.Module] = [] + for _ in range(3): + stage.append(block(block_setting[0].input_channels, layer_scale, 0)) + self.stage_smooth = nn.Sequential(*stage).to(device) + + self.first = nn.Sequential( + nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2), + norm_layer(block_setting[0].out_channels), + ).to(device) + + + total_stage_blocks = sum(cnf.num_layers for cnf in block_setting) + stage_block_id = 0 + layers_list: List[nn.Module] = [] + + for cnf in block_setting[1:]: + if cnf.input_channels is not None: + # Bottlenecks + stage: List[nn.Module] = [] + for _ in range(cnf.num_layers): + # adjust stochastic depth probability based on the depth of the stage block + sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0) + stage.append(block(cnf.input_channels, layer_scale, sd_prob)) + stage_block_id += 1 + layers_list.append(nn.Sequential(*stage).to(device)) + else: + layers_list.append(nn.Identity().to(device)) + + if cnf.out_channels is not None: + # Upsampling + layers_list.append( + nn.Sequential( + norm_layer(cnf.input_channels), + nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2), + ).to(device) + ) + + self.layers = nn.Sequential(*layers_list[:-1]).to(device) + self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device) + + stage: List[nn.Module] = [] + for _ in range(cnf.num_layers): + stage.append(block(block_setting[-1].input_channels, layer_scale, 0)) + self.stage_mask = nn.Sequential(*stage).to(device) + + self.last = nn.Conv2d( + block_setting[-1].input_channels, + 3, + kernel_size=3, + padding='same', + ).to(device) + + self.last_mask = nn.Conv2d( + block_setting[-1].input_channels, + 1, + kernel_size=3, + padding='same', + ).to(device) + + for m in self.modules(): + if isinstance(m, (nn.Conv2d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + + + def _forward_impl(self, x) -> Tensor: + if self.use_skip: + raise NotImplementedError + input, features = x + else: + input = x + x = self.stage_smooth(input) + x = self.first(x) + f = self.layers(x) + x = self.stage_pose(f) + pose = torch.sigmoid(self.last(x)) + + x = self.stage_mask(f) + mask = torch.sigmoid(self.last_mask(x)) + pose = pose * (mask > 0.5).to(torch.float32) + + return pose, mask + + def forward(self, x) -> Tensor: + return self._forward_impl(x) + + +class PoseAE(nn.Module): + def __init__(self, features=256, device='cuda'): + super().__init__() + + + block_setting = [ + CNBlockConfig(96, 192, 3), + CNBlockConfig(192, 384, 3), + CNBlockConfig(384, features, 9), + CNBlockConfig(features, None, 3), + ] + + block_setting_reverse = [ + CNBlockConfig(features, 384, 3), + CNBlockConfig(384, 192, 9), + CNBlockConfig(192, 96, 3), + CNBlockConfig(96, 27, 3), + CNBlockConfig(27, 9, 3), + CNBlockConfig(9, None, 3), + ] + + self.encoder = ConvNeXt(block_setting, input_features=4, device=device) + self.decoder = ConvNeXtReversePose(block_setting_reverse, device=device) + + self.flatten = nn.Flatten() + + + def forward(self, input) -> torch.Tensor: + mask, pose_in = input + x, _ = self.encoder(torch.cat([mask, pose_in], dim=1)) + + pose_out, mask_out = self.decoder(x) + + return pose_out, mask_out + + +if __name__ == '__main__': + # test the model + model = PoseAE() + # model = torch.compile(model) + print(model) + print(sum(p.numel() for p in model.parameters() if p.requires_grad)) + print(sum(p.numel() for p in model.parameters())) + + x = torch.rand((4, 1, 512, 512)).to('cuda') + y = torch.rand((4, 3, 512, 512)).to('cuda') + output = model((x, y)) + + + + pass \ No newline at end of file diff --git a/scripts/pose_ae_train.py b/scripts/pose_ae_train.py new file mode 100644 index 0000000..40276f0 --- /dev/null +++ b/scripts/pose_ae_train.py @@ -0,0 +1,182 @@ +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"] = '7' +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 models.pose_ae import PoseAE +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_ae_2', notes='now with random poses', dir='/data-local/127/temp', tags=['pose']) + artifact = wandb.Artifact('code', type='code') + artifact.add_file('scripts/pose_ae_train.py') + artifact.add_file('models/pose_ae.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 = model([batch['mask'], batch['pose']]) + + mse = loss_fn(pred_pose, batch['pose']) + mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask'].squeeze(1)) + + loss: torch.Tensor = mse + mse_mask + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + + return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask } + + +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']]) + + mask = torch.permute(batch['mask'], (0, 2, 3, 1)) + pose = torch.permute(batch['pose'], (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, pred_mask), axis=2).squeeze(-1).cpu().numpy() + poses = torch.concat((pose, 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) + 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" + + 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=True) + valid_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=True) + 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 = PoseAE(device=device).to(device) + + loss_fn = 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) + +# 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