diff options
Diffstat (limited to 'scripts/validation.py')
| -rw-r--r-- | scripts/validation.py | 261 |
1 files changed, 261 insertions, 0 deletions
diff --git a/scripts/validation.py b/scripts/validation.py new file mode 100644 index 0000000..5b1638d --- /dev/null +++ b/scripts/validation.py @@ -0,0 +1,261 @@ +import torch as th +from torch import nn +from torch.utils.data import DataLoader +import numpy as np +from einops import rearrange, repeat, reduce +from scripts.utils.configuration import Configuration +from model.loci import Loci +import time +import lpips +from skimage.metrics import structural_similarity as ssimloss +from skimage.metrics import peak_signal_noise_ratio as psnrloss + +def validation_adept(valloader: DataLoader, net: Loci, cfg: Configuration, device): + + # memory + mseloss = nn.MSELoss() + avgloss = 0 + start_time = time.time() + + with th.no_grad(): + for i, input in enumerate(valloader): + + # get input frame and target frame + tensor = input[0].float().to(device) + background_fix = input[1].to(device) + + # apply skip frames + tensor = tensor[:,range(0, tensor.shape[1], cfg.defaults.skip_frames)] + sequence_len = tensor.shape[1] + + # initial frame + input = tensor[:,0] + target = th.clip(tensor[:,0], 0, 1) + error_last = None + + # placehodlers + mask_cur = None + mask_last = None + rawmask_last = None + position_last = None + gestalt_last = None + priority_last = None + gt_positions_target = None + slots_occlusionfactor = None + + # loop through frames + for t_index,t in enumerate(range(-cfg.defaults.teacher_forcing, sequence_len-1)): + + # move to next frame + t_run = max(t, 0) + input = tensor[:,t_run] + target = th.clip(tensor[:,t_run+1], 0, 1) + + # obtain prediction + ( + output_next, + position_next, + gestalt_next, + priority_next, + mask_next, + rawmask_next, + object_next, + background, + slots_occlusionfactor, + output_cur, + position_cur, + gestalt_cur, + priority_cur, + mask_cur, + rawmask_cur, + object_cur, + position_encoder_cur, + slots_bounded, + slots_partially_occluded_cur, + slots_occluded_cur, + slots_partially_occluded_next, + slots_occluded_next, + slots_closed, + output_hidden, + largest_object, + rawmask_hidden, + object_hidden + ) = net( + input, + error_last, + mask_last, + rawmask_last, + position_last, + gestalt_last, + priority_last, + background_fix, + slots_occlusionfactor, + reset = (t == -cfg.defaults.teacher_forcing), + evaluate=True, + warmup = (t < 0), + shuffleslots = False, + reset_mask = (t <= 0), + allow_spawn = True, + show_hidden = False, + clean_slots = (t <= 0), + ) + + # 1. Track error + if t >= 0: + loss = mseloss(output_next, target) + avgloss += loss.item() + + # 2. Remember output + mask_last = mask_next.clone() + rawmask_last = rawmask_next.clone() + position_last = position_next.clone() + gestalt_last = gestalt_next.clone() + priority_last = priority_next.clone() + + # 3. Error for next frame + # background error + bg_error_cur = th.sqrt(reduce((input - background)**2, 'b c h w -> b 1 h w', 'mean')).detach() + bg_error_next = th.sqrt(reduce((target - background)**2, 'b c h w -> b 1 h w', 'mean')).detach() + + # prediction error + error_next = th.sqrt(reduce((target - output_next)**2, 'b c h w -> b 1 h w', 'mean')).detach() + error_next = th.sqrt(error_next) * bg_error_next + error_last = error_next.clone() + + print(f"Validation loss: {avgloss / len(valloader.dataset):.2e}, Time: {time.time() - start_time}") + + pass + + +def validation_clevrer(valloader: DataLoader, net: Loci, cfg: Configuration, device): + + # memory + mseloss = nn.MSELoss() + lpipsloss = lpips.LPIPS(net='vgg').to(device) + avgloss_mse = 0 + avgloss_lpips = 0 + avgloss_psnr = 0 + avgloss_ssim = 0 + start_time = time.time() + + burn_in_length = 6 + rollout_length = 42 + + with th.no_grad(): + for i, input in enumerate(valloader): + + # get input frame and target frame + tensor = input[0].float().to(device) + background_fix = input[1].to(device) + + # apply skip frames + tensor = tensor[:,range(0, tensor.shape[1], cfg.defaults.skip_frames)] + sequence_len = tensor.shape[1] + + # initial frame + input = tensor[:,0] + target = th.clip(tensor[:,0], 0, 1) + error_last = None + + # placehodlers + mask_cur = None + mask_last = None + rawmask_last = None + position_last = None + gestalt_last = None + priority_last = None + gt_positions_target = None + slots_occlusionfactor = None + + # loop through frames + for t_index,t in enumerate(range(-cfg.defaults.teacher_forcing, min(burn_in_length + rollout_length-1, sequence_len-1))): + + # move to next frame + t_run = max(t, 0) + input = tensor[:,t_run] + target = th.clip(tensor[:,t_run+1], 0, 1) + if t_run >= burn_in_length: + blackout = th.tensor((np.random.rand(valloader.batch_size) < 0.2)[:,None,None,None]).float().to(device) + input = blackout * (input * 0) + (1-blackout) * input + error_last = blackout * (error_last * 0) + (1-blackout) * error_last + + # obtain prediction + ( + output_next, + position_next, + gestalt_next, + priority_next, + mask_next, + rawmask_next, + object_next, + background, + slots_occlusionfactor, + output_cur, + position_cur, + gestalt_cur, + priority_cur, + mask_cur, + rawmask_cur, + object_cur, + position_encoder_cur, + slots_bounded, + slots_partially_occluded_cur, + slots_occluded_cur, + slots_partially_occluded_next, + slots_occluded_next, + slots_closed, + output_hidden, + largest_object, + rawmask_hidden, + object_hidden + ) = net( + input, + error_last, + mask_last, + rawmask_last, + position_last, + gestalt_last, + priority_last, + background_fix, + slots_occlusionfactor, + reset = (t == -cfg.defaults.teacher_forcing), + evaluate=True, + warmup = (t < 0), + shuffleslots = False, + reset_mask = (t <= 0), + allow_spawn = True, + show_hidden = False, + clean_slots = (t <= 0), + ) + + # 1. Track error + if t >= 0: + loss_mse = mseloss(output_next, target) + loss_ssim = np.sum([ssimloss(output_next[i].cpu().numpy(), target[i].cpu().numpy(), channel_axis=0,gaussian_weights=True,sigma=1.5,use_sample_covariance=False,data_range=1) for i in range(output_next.shape[0])]), + loss_psnr = np.sum([psnrloss(output_next[i].cpu().numpy(), target[i].cpu().numpy(), data_range=1) for i in range(output_next.shape[0])]), + loss_lpips = th.sum(lpipsloss(output_next*2-1, target*2-1)) + + avgloss_mse += loss_mse.item() + avgloss_ssim += loss_ssim[0].item() + avgloss_psnr += loss_psnr[0].item() + avgloss_lpips += loss_lpips.item() + + # 2. Remember output + mask_last = mask_next.clone() + rawmask_last = rawmask_next.clone() + position_last = position_next.clone() + gestalt_last = gestalt_next.clone() + priority_last = priority_next.clone() + + # 3. Error for next frame + # background error + bg_error_next = th.sqrt(reduce((target - background)**2, 'b c h w -> b 1 h w', 'mean')).detach() + + # prediction error + error_next = th.sqrt(reduce((target - output_next)**2, 'b c h w -> b 1 h w', 'mean')).detach() + error_next = th.sqrt(error_next) * bg_error_next + error_last = error_next.clone() + + print(f"MSE loss: {avgloss_mse / len(valloader.dataset):.2e}, LPIPS loss: {avgloss_lpips / len(valloader.dataset):.2e}, PSNR loss: {avgloss_psnr / len(valloader.dataset):.2e}, SSIM loss: {avgloss_ssim / len(valloader.dataset):.2e}, Time: {time.time() - start_time}") + + pass
\ No newline at end of file |