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

use new colors for pose

Browse Source
This commit is contained in:
luni
2023-06-08 10:54:34 +00:00
parent fdc876025a
commit 882ddf5578
3 changed files with 23 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/pose_encoder_train.py
+15 -17
View File
@@ -94,19 +94,20 @@ def valid_step(batch):
symm_rotation = loss_fn_rotation(pred_rotation)
loss = mse + mse_mask + symm_rotation
return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
def plot_everything(batch, type: str):
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
model.eval()
with torch.inference_mode():
pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_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_out = torch.permute(batch['pose_out'], (0, 2, 3, 1))
pred_mask = torch.permute(pred_mask, (0, 2, 3, 1)).detach()
pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)).detach()
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)
@@ -121,7 +122,7 @@ if __name__ == '__main__':
device = "cuda" if torch.cuda.is_available() else "cpu"
batch_size = 1
batch_size = 4
use_cache = False
train_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
@@ -156,8 +157,6 @@ if __name__ == '__main__':
# if iteration == 0:
# to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if not args['debug']:
wandb.log(to_log, step=iteration)
train_accumulator.reset()
iteration += 1
# print(f'Iteration {iteration} finished')
@@ -165,18 +164,17 @@ if __name__ == '__main__':
epoch += 1
print(f'Epoch {epoch} finished')
if not iteration % 1000 == 0:
continue
# if not iteration % 1000 == 0:
# continue
for batch in valid_dataloader:
valid_loss = valid_step(batch)
valid_accumulator.update(valid_loss)
# 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()
# 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)