add debug logging to cvae_train

configs/abl_all.py

@@ -2,12 +2,13 @@ from utils.convnext_utils import CNBlockConfig
 
 config = dict()
 
-config['gpu'] = '0'
-config['name'] = 'abl_torch_all'
+config['gpu'] = '6'
+config['name'] = 'abl_torch_all_2'
 config['notes'] = ''
 config['tags'] = ['iccv', 'abl_all', 'torch']
 config['script'] = 'scripts/cvae.py'
 
+config['pretrained_model'] = 'checkpoints/iccv/abl_torch_all'
 
 config['use_skinning'] =              True # and deshaping
 config['remove_global_orientation'] = True
@@ -43,15 +44,12 @@ config['static_model'] = 'checkpoints/iccv/static_torch_orig_2'
 config['body_model'] = 'checkpoints/iccv/body_torch_1'
 config['curvature_model'] = 'checkpoints/iccv/curv_torch_1'
 
-
 config['smpl_res'] = 14475
 if config['is_cape']:
   config['smpl_res'] = 3638
 
 config['latent_dim'] = 128
 
-config['pretrained_model'] = ''
-
 ###############    
 config['lr'] = 0.0001#                 0.0001
 ###############

configs/ae_curv_1.py

@@ -53,7 +53,6 @@ block_setting_reverse = [
     CNBlockConfig(64, 32, 9),
     CNBlockConfig(32, 16, 3),
     CNBlockConfig(16, 3, 3),
-    # CNBlockConfig(3, 3, 3),
 ]
 
 config['cond_ae']['encoder_block_settings'] = block_setting

configs/ae_static_orig.py

@@ -7,6 +7,8 @@ config['name'] = 'static_torch_orig_1'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_orig_1'
+
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/train'
 config['valid_dir'] = '/data-local/127/valid'
@@ -60,7 +62,6 @@ config['cond_ae']['activation'] = 'tanh'
 
 config['lr'] = 0.0005
 config['cond_ae']['upscale_output'] = False
-config['pretrained_model'] = ''
 
 config['signature'] = { 'id',
                         'frame_nr',

configs/ae_static_orig_2.py

@@ -7,6 +7,8 @@ config['name'] = 'static_torch_orig_2'
 config['notes'] = ' ' 
 config['tags'] = ['torch', 'static'] 
 
+config['pretrained_model'] = 'checkpoints/iccv/static_torch_orig_2'
+
 config['batch_size'] = 4
 config['train_dir'] = '/data-local/127/train'
 config['valid_dir'] = '/data-local/127/valid'

docker/Dockerfile2

@@ -0,0 +1,46 @@
+# Use an existing image as the base
+FROM nvcr.io/nvidia/pytorch:23.04-py3
+# FROM pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel
+# FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
+# 
+# Add the current user to the image
+ARG USER_NAME=lhunor
+ARG USER_ID=11017
+ARG GROUP_NAME=guest
+ARG GROUP_ID=11000
+
+# Update the system and install required packages
+RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git apt-utils htop
+
+# Create a new group with the same GID as the host
+RUN groupadd --gid $GROUP_ID $GROUP_NAME
+
+# Create a new user with the same UID and GID as the host
+RUN useradd --uid $USER_ID --gid $GROUP_ID --shell /bin/bash --create-home $USER_NAME
+
+# don't require password with sudo, for convenience
+# not the safest thing to do, but hopefully okay inside the container
+RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
+
+# Add the new user to the sudo group
+RUN usermod -aG sudo $USER_NAME
+
+# Set the default user to the new user
+USER $USER_NAME
+
+WORKDIR /workspace
+
+ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
+RUN pip install --upgrade pip
+COPY requirements.txt /tmp/
+RUN pip install --requirement /tmp/requirements.txt
+RUN pip install torch-cluster -f https://data.pyg.org/whl/torch-2.0.0+cu117.html
+
+
+# defining the volumes that will be mounted
+VOLUME [ "/data-net" ]
+VOLUME [ "/data-local" ]
+VOLUME [ "/home-net" ]
+
+# Start the container with a bash shell
+CMD ["/bin/bash"]

docker/docker-compose.yml

@@ -19,3 +19,24 @@ services:
         reservations:
           devices:
             - capabilities: [gpu]
+
+
+  lhunor_torch_2:
+    image: lhunor_torch:nv23.04
+    container_name: lhunor_torch_nv23.04
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      - '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]

docker/requirements.txt

@@ -1,3 +1,4 @@
+torch==2.0.0    
 numpy
 matplotlib
 opencv-python

run.py

@@ -1,22 +1,24 @@
 import os
 import sys
-from configs.ae_curv_1 import config
+from configs.abl_all import config
 import argparse
 
-parser = argparse.ArgumentParser(description='')
-parser.add_argument('--debug', action='store_true')
-args = vars(parser.parse_args())
+# parser = argparse.ArgumentParser(description='')
+# parser.add_argument('--debug', action='store_true')
+# args = vars(parser.parse_args())
 
 
 
 if __name__ == '__main__':
-    # import wandb
-    # wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
+    import wandb
+    # wandb.login(key='******************************************')
+    wandb.login(host='https://wandb.lnrweb.com', key='local-40a577c9eccab708f1aa81027f21a995596a5838')
     name = config['name']
     script = config['script']
     config = [key for key in sys.modules.keys()  if key.startswith('configs.')][-1]
     current_file = '/'.join(__file__.split('/')[-2:])
     
+    print(f'nohup python {script} --name lhunor_{name} --config {config} > nohup/torch/{name}.out 2>&1 &')
     if args['debug']:
         os.system(f'python {script} --name lhunor_{name} --config {config}')
     else:

scripts/ae_static_train.py

@@ -2,7 +2,16 @@ import os
 import sys
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-from configs.ae_static import config
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 
@@ -98,6 +107,8 @@ def valid_step(batch):
        
 
 def plot_everything(batch: dict, type: str):
+    model.eval()
+    with torch.inference_mode():
         pred = model([batch['uv_static'], batch['uv_mask']]).detach()
     
     mesh_pred = uv2mesh(pred)
@@ -162,9 +173,9 @@ if __name__ == '__main__':
     init_wandb()
     # device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    parser = argparse.ArgumentParser(description='')
-    parser.add_argument('--name', type=str)
-    args = vars(parser.parse_args())
+    # parser = argparse.ArgumentParser(description='')
+    # parser.add_argument('--name', type=str)
+    # args = vars(parser.parse_args())
     
     resolution_modifier = ''
     if config['resolution'] != 512:

scripts/cvae.py

@@ -7,6 +7,7 @@ import argparse
 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'])
@@ -128,8 +129,6 @@ def train_step(batch):
             }
  
 
-
-
 def valid_step(batch):
     if config['use_skinning']:
         mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
@@ -152,7 +151,6 @@ def valid_step(batch):
         else:
             mesh_pred_posed = mesh_pred
 
-
         e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
         c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
         c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
@@ -221,6 +219,19 @@ def plot_everything(batch: dict, type: str):
 
     gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
 
+    # # plot garment_unposed, gt and mesh_posed next to eachother in plotly subplots
+    # from plotly.subplots import make_subplots
+    # fig = make_subplots(rows=1, cols=4, specs=[[{'type': 'surface'}, {'type': 'surface'}, {'type': 'surface'}, {'type': 'surface'}]])
+    # fig.add_trace(go.Mesh3d(x=gt[0][:,0], y=gt[0][:,2], z=gt[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='red', opacity=0.5), row=1, col=1)
+    # fig.add_trace(go.Mesh3d(x=garment_unposed[0][:,0], y=garment_unposed[0][:,2], z=garment_unposed[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='blue', opacity=0.5), row=1, col=2)
+    # # add new subplot for gt_posed
+    # fig.add_trace(go.Mesh3d(x=gt_posed[0][:,0], y=gt_posed[0][:,2], z=gt_posed[0][:,1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='red', opacity=0.5), row=1, col=3)
+    # fig.add_trace(go.Mesh3d(x=mesh_posed[0][mask[0],0], y=mesh_posed[0][mask[0],2], z=mesh_posed[0][mask[0],1], i=faces[0][:,0], j=faces[0][:,1], k=faces[0][:,2], color='green', opacity=0.5), row=1, col=4)
+    # # make each subplot between -1 and 1
+    # fig.update_layout(scene_aspectmode='cube', scene2_aspectmode='cube', scene3_aspectmode='cube', scene4_aspectmode='cube')
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene2=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene3=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])), scene4=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+
+    # fig.show()
 
     mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy()
     body = batch['mesh_body_posed'].cpu().numpy()
@@ -412,9 +423,10 @@ def test():
 
 
 def train():
+    global g_balancing_coeff
+    global step
     valid_iterator = iter(valid_dataloader)
 
-    step = 0
     for batch in train_dataloader:
         train_loss = train_step(batch)
         train_accumulator.update(train_loss)
@@ -433,15 +445,14 @@ def train():
             if step == 0:
                 to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
+            if not args['debug']:
                 wandb.log(to_log, step=step)
             train_accumulator.reset()
             valid_accumulator.reset()
             print(step)
 
         if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
-            if balancing_activated_step == 0:
-                balancing_activated_step = step
-            g_balancing_coeff = g_balancing_coeff + 1 / 10000 
+            g_balancing_coeff += 1 / 10000 
 
         if step % 10000 == 0:
             ckpt_saver(model, step, optimizer)
@@ -450,11 +461,8 @@ def train():
 
 
 if __name__ == '__main__':
+    if not args['debug']:
         init_wandb()
-    # parser = argparse.ArgumentParser(description='')
-    # parser.add_argument('--name', type=str)
-    # parser.add_argument('--debug', action='store_true')
-    # args = vars(parser.parse_args())
 
     smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
@@ -468,14 +476,19 @@ if __name__ == '__main__':
 
     model = CVAE(config)
 
+
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    balancing_activated_step = 0
     g_balancing_coeff = config['starting_balancing_coeff']
 
     ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
 
+    step = 0
     if config['pretrained_model'] != '':
-        load_model(model, config['pretrained_model'])
+        step = load_model(model, config['pretrained_model'])
+
+    if not args['debug']:
+        model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
         
     train()

smpl/smpl_torch.py

@@ -204,7 +204,7 @@ class SMPL(Module):
         )
       )
     
-    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+    # joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
             
     T = torch.stack([torch.tensordot(results[i], self.weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
     rest_shape_h = torch.cat(
@@ -214,11 +214,18 @@ class SMPL(Module):
     v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
     result = v
 
-    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
-    result = result - root_joint
-    joints = joints - root_joint
+    # root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    # result = result - root_joint
+    # joints = joints - root_joint
+
+    # # plut result as a plotly 3d mesh
+    # import plotly.graph_objects as go
+    # fig = go.Figure(data=[go.Mesh3d(x=result[0,:,0].cpu().numpy(), y=result[0,:,2].cpu().numpy(), z=result[0,:,1].cpu().numpy(), color='lightpink', opacity=0.50, i=self.f[:,0], j=self.f[:,1], k=self.f[:,2], name='body')])
+    # # plot the joints as spheres
+    # fig.add_trace(go.Scatter3d(x=joints[0,:,0].cpu().numpy(), y=joints[0,:,2].cpu().numpy(), z=joints[0,:,1].cpu().numpy(), mode='markers', marker=dict(size=5, color='red'), name='joints'))
+    # fig.show()
 
     if remove and verts is None:
       result = result[:,self.to_keep,:]
 
-    return result, joints
+    return result, None # joints

utils/checkpoints.py

@@ -35,6 +35,8 @@ def load_model(model, checkpoint, optimizer=None, iteration=''):
     # list of checkpoint files in the folder checkpoints
     checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
     checkpoint = torch.load(checkpoints[-1])
+    iteration = checkpoints[-1].split('_')[-1].split('.')[0]
     model.load_state_dict(checkpoint['model'])
     if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'])
+    return int(iteration) if iteration != '' else 0