From 0b9daf4841c980d6471812736878cf1c4daf545e Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 17 Jun 2023 09:33:23 +0000 Subject: [PATCH] add checkpoint saving --- scripts/pose_ae_train.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/scripts/pose_ae_train.py b/scripts/pose_ae_train.py index 40276f0..7bf43b5 100644 --- a/scripts/pose_ae_train.py +++ b/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))