feature manipulator pipeline

2023-06-17 16:12:15 +00:00
parent 37c67d5cab
commit d9021672e4
2 changed files with 278 additions and 0 deletions
@@ -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