add checkpoint saving

scripts/pose_ae_train.py

@@ -32,6 +32,7 @@ from torch.profiler import profile, record_function, ProfilerActivity
 
 from data.pose_dataset import PoseDataset
 from utils.utils import Accumulator
+from utils.checkpoints import CheckpointSaver, load_model
 from models.pose_ae import PoseAE
 from models.convnext import ConvNeXt
 import plotly.express as ex
@@ -40,8 +41,10 @@ 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_ae_2'
+
 def init_wandb():
-    run = wandb.init(project='reconstruction', name='pose_ae_2', notes='now with random poses', dir='/data-local/127/temp', tags=['pose'])
+    run = wandb.init(project='reconstruction', name=name, 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')
@@ -125,9 +128,14 @@ def train():
                 # 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')
 
+            if iteration % 10000 == 0:
+                ckpt_saver_enc(model.encoder, iteration, optimizer, scaler=None)
+                ckpt_saver_dec(model.decoder, iteration, optimizer, scaler=None)
+
+            iteration += 1
+            
         epoch += 1
         print(f'Epoch {epoch} finished')
 
@@ -172,9 +180,17 @@ if __name__ == '__main__':
     # lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
     # optimizer = optim.SGD(model.parameters(), lr=0.1)
 
+    ckpt_saver_enc = CheckpointSaver('checkpoints/reconstruction/' + name + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/reconstruction/' + name + '_dec', 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))