pose autoencoder pipeline

data/pose_dataset.py

@@ -17,10 +17,11 @@ from data.pose_renderer import PoseRenderer
 import trimesh
 
 
-class Dataset(Dataset):
+class PoseDataset(Dataset):
-    def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True):
+    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,29 +69,39 @@ 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()
 
+        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 = 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]))
+        # 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 = torch.tensor(mesh.vertices, device=body[0].device).unsqueeze(0).to(torch.float32)
-        body_back[0,:,0] = -body_back[0,:,0] 
+        # 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_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_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
@@ -98,11 +109,13 @@ 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_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_back = render_back * mask_back
+        if not self.ae_mode:
             render_out = render_out * mask_out
 
         # print(f'render took {time() - start} seconds')
@@ -110,16 +123,20 @@ class Dataset(Dataset):
         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
         return 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:

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)

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

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