luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

add checkpoint saving

Browse Source
This commit is contained in:
luni
2023-06-17 09:33:23 +00:00
parent 2c873b255e
commit 0b9daf4841
1 changed files with 18 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/pose_ae_train.py
+18 -2
View File
@@ -32,6 +32,7 @@ from torch.profiler import profile, record_function, ProfilerActivity
from data.pose_dataset import PoseDataset from data.pose_dataset import PoseDataset
from utils.utils import Accumulator from utils.utils import Accumulator
from utils.checkpoints import CheckpointSaver, load_model
from models.pose_ae import PoseAE from models.pose_ae import PoseAE
from models.convnext import ConvNeXt from models.convnext import ConvNeXt
import plotly.express as ex 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) ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
name = 'pose_ae_2'
def init_wandb(): 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 = wandb.Artifact('code', type='code')
artifact.add_file('scripts/pose_ae_train.py') artifact.add_file('scripts/pose_ae_train.py')
artifact.add_file('models/pose_ae.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 # figure out when to log when using validation set as well
if not args['debug']: if not args['debug']:
wandb.log(to_log, step=iteration) wandb.log(to_log, step=iteration)
iteration += 1
# print(f'Iteration {iteration} finished') # 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 epoch += 1
print(f'Epoch {epoch} finished') 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) # lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
# optimizer = optim.SGD(model.parameters(), lr=0.1) # 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 # Use the profiler to track memory usage
# with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
# with record_function("model_inference"): # 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() train()
# print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10)) # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))