pose pipeline trial
This commit is contained in:
@@ -26,6 +26,7 @@ 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
|
||||
|
||||
@@ -40,7 +41,7 @@ ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='reconstruction', name='pose_1', notes='', dir='/data-local/127/temp', tags=['pose'])
|
||||
run = wandb.init(project='reconstruction', name='pose_2', notes='have the back of the body also as the main input', dir='/data-local/127/temp', tags=['pose'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('scripts/pose_encoder_train.py')
|
||||
artifact.add_file('models/pose_encoder.py')
|
||||
@@ -69,30 +70,34 @@ class SymmetricLoss(nn.Module):
|
||||
def train_step(batch):
|
||||
model.train()
|
||||
|
||||
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
|
||||
# pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
|
||||
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 = 0 #loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
|
||||
symm_rotation = 0 #loss_fn_rotation(pred_rotation)
|
||||
|
||||
loss = mse #+ mse_mask + symm_rotation
|
||||
print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}')
|
||||
# print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}')
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
with torch.inference_mode():
|
||||
mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
|
||||
symm_rotation = loss_fn_rotation(pred_rotation)
|
||||
|
||||
return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask, 'train/symm_rotation': symm_rotation }
|
||||
|
||||
|
||||
def valid_step(batch):
|
||||
model.eval()
|
||||
with torch.inference_mode():
|
||||
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
|
||||
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, batch['mask_out'])
|
||||
mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
|
||||
symm_rotation = loss_fn_rotation(pred_rotation)
|
||||
loss = mse + mse_mask + symm_rotation
|
||||
loss = mse
|
||||
|
||||
return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
|
||||
|
||||
@@ -100,24 +105,76 @@ def valid_step(batch):
|
||||
def plot_everything(batch, type: str):
|
||||
model.eval()
|
||||
with torch.inference_mode():
|
||||
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
|
||||
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
|
||||
|
||||
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_back = torch.permute(batch['pose_back'], (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)
|
||||
poses = torch.concat((pose, pose_out, pred_pose), axis=2)
|
||||
masks_imgs = [wandb.Image((masks[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
|
||||
poses_imgs = [wandb.Image((poses[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
|
||||
return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, } }
|
||||
masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy()
|
||||
poses = torch.concat((pose, pose_back, 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])]
|
||||
|
||||
# do a batchwise min-max normalization of pred_rotation
|
||||
pred_rotation_img = pred_rotation.cpu().numpy()
|
||||
pred_rotation_img = (pred_rotation_img - pred_rotation_img.min(axis=(1, 2,), keepdims=True)) / (pred_rotation_img.max(axis=(1, 2), keepdims=True) - pred_rotation_img.min(axis=(1, 2), keepdims=True))
|
||||
pred_rotation_imgs = [wandb.Image(pred_rotation_img[i]) for i in range(pred_pose.shape[0])]
|
||||
|
||||
return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, f'{type}/rotation': pred_rotation_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)
|
||||
iteration += 1
|
||||
# print(f'Iteration {iteration} finished')
|
||||
|
||||
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"
|
||||
@@ -141,45 +198,11 @@ if __name__ == '__main__':
|
||||
lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
|
||||
# optimizer = optim.SGD(model.parameters(), lr=0.1)
|
||||
|
||||
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 < 1000 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()
|
||||
iteration += 1
|
||||
# print(f'Iteration {iteration} finished')
|
||||
|
||||
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)
|
||||
|
||||
|
||||
# Use the profiler to track memory usage
|
||||
# with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
|
||||
# with record_function("model_inference"):
|
||||
train()
|
||||
|
||||
# print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
|
||||
|
||||
pass
|
||||
Reference in New Issue
Block a user