procgne/Plonk
0
1from typing import Any2import pytorch_lightning as L3import torch4import torch.nn as nn5from hydra.utils import instantiate6import copy7import pandas as pd8import numpy as np9from tqdm import tqdm10from utils.manifolds import Sphere11from torch.func import jacrev, vjp, vmap12from torchdiffeq import odeint13from geoopt import ProductManifold, Euclidean14from models.samplers.riemannian_flow_sampler import ode_riemannian_flow_sampler15 16 17class DiffGeolocalizer(L.LightningModule):18 def __init__(self, cfg):19 super().__init__()20 self.cfg = cfg21 self.network = instantiate(cfg.network)22 # self.network = torch.compile(self.network, fullgraph=True)23 self.input_dim = cfg.network.input_dim24 self.train_noise_scheduler = instantiate(cfg.train_noise_scheduler)25 self.inference_noise_scheduler = instantiate(cfg.inference_noise_scheduler)26 self.data_preprocessing = instantiate(cfg.data_preprocessing)27 self.cond_preprocessing = instantiate(cfg.cond_preprocessing)28 self.preconditioning = instantiate(cfg.preconditioning)29 30 self.ema_network = copy.deepcopy(self.network).requires_grad_(False)31 self.ema_network.eval()32 self.postprocessing = instantiate(cfg.postprocessing)33 self.val_sampler = instantiate(cfg.val_sampler)34 self.test_sampler = instantiate(cfg.test_sampler)35 self.loss = instantiate(cfg.loss)(36 self.train_noise_scheduler,37 )38 self.val_metrics = instantiate(cfg.val_metrics)39 self.test_metrics = instantiate(cfg.test_metrics)40 self.manifold = instantiate(cfg.manifold) if hasattr(cfg, "manifold") else None41 42 self.interpolant = cfg.interpolant43 44 def training_step(self, batch, batch_idx):45 with torch.no_grad():46 batch = self.data_preprocessing(batch)47 batch = self.cond_preprocessing(batch)48 batch_size = batch["x_0"].shape[0]49 loss = self.loss(self.preconditioning, self.network, batch).mean()50 self.log(51 "train/loss",52 loss,53 sync_dist=True,54 on_step=True,55 on_epoch=True,56 batch_size=batch_size,57 )58 return loss59 60 def on_before_optimizer_step(self, optimizer):61 if self.global_step == 0:62 no_grad = []63 for name, param in self.network.named_parameters():64 if param.grad is None:65 no_grad.append(name)66 if len(no_grad) > 0:67 print("Parameters without grad:")68 print(no_grad)69 70 def on_validation_start(self):71 self.validation_generator = torch.Generator(device=self.device).manual_seed(72 340773 )74 self.validation_generator_ema = torch.Generator(device=self.device).manual_seed(75 340776 )77 78 def validation_step(self, batch, batch_idx):79 batch = self.data_preprocessing(batch)80 batch = self.cond_preprocessing(batch)81 batch_size = batch["x_0"].shape[0]82 loss = self.loss(83 self.preconditioning,84 self.network,85 batch,86 generator=self.validation_generator,87 ).mean()88 self.log(89 "val/loss",90 loss,91 sync_dist=True,92 on_step=False,93 on_epoch=True,94 batch_size=batch_size,95 )96 if hasattr(self, "ema_model"):97 loss_ema = self.loss(98 self.preconditioning,99 self.ema_network,100 batch,101 generator=self.validation_generator_ema,102 ).mean()103 self.log(104 "val/loss_ema",105 loss_ema,106 sync_dist=True,107 on_step=False,108 on_epoch=True,109 batch_size=batch_size,110 )111 # nll = -self.compute_exact_loglikelihood(batch).mean()112 # self.log(113 # "val/nll",114 # nll,115 # sync_dist=True,116 # on_step=False,117 # on_epoch=True,118 # batch_size=batch_size,119 # )120 121 # def on_validation_epoch_end(self):122 # metrics = self.val_metrics.compute()123 # for metric_name, metric_value in metrics.items():124 # self.log(125 # f"val/{metric_name}",126 # metric_value,127 # sync_dist=True,128 # on_step=False,129 # on_epoch=True,130 # )131 132 def on_test_start(self):133 self.test_generator = torch.Generator(device=self.device).manual_seed(3407)134 135 def test_step_simple(self, batch, batch_idx):136 batch = self.data_preprocessing(batch)137 batch = self.cond_preprocessing(batch)138 batch_size = batch["x_0"].shape[0]139 if isinstance(self.manifold, Sphere):140 x_N = self.manifold.random_base(141 batch_size,142 self.input_dim,143 device=self.device,144 )145 x_N = x_N.reshape(batch_size, self.input_dim)146 else:147 x_N = torch.randn(148 batch_size,149 self.input_dim,150 device=self.device,151 generator=self.test_generator,152 )153 cond = batch[self.cfg.cond_preprocessing.output_key]154 155 samples = self.sample(156 x_N=x_N,157 cond=cond,158 stage="val",159 generator=self.test_generator,160 cfg=self.cfg.cfg_rate,161 )162 self.test_metrics.update({"gps": samples}, batch)163 if self.cfg.compute_nll:164 nll = -self.compute_exact_loglikelihood(batch, cfg=0).mean()165 self.log(166 "test/NLL",167 nll,168 sync_dist=True,169 on_step=False,170 on_epoch=True,171 batch_size=batch_size,172 )173 174 def test_best_nll(self, batch, batch_idx):175 batch = self.data_preprocessing(batch)176 batch = self.cond_preprocessing(batch)177 batch_size = batch["x_0"].shape[0]178 num_sample_per_cond = 32179 if isinstance(self.manifold, Sphere):180 x_N = self.manifold.random_base(181 batch_size * num_sample_per_cond,182 self.input_dim,183 device=self.device,184 )185 x_N = x_N.reshape(batch_size * num_sample_per_cond, self.input_dim)186 else:187 x_N = torch.randn(188 batch_size * num_sample_per_cond,189 self.input_dim,190 device=self.device,191 generator=self.test_generator,192 )193 cond = (194 batch[self.cfg.cond_preprocessing.output_key]195 .unsqueeze(1)196 .repeat(1, num_sample_per_cond, 1)197 .view(-1, batch[self.cfg.cond_preprocessing.output_key].shape[-1])198 )199 samples = self.sample_distribution(200 x_N,201 cond,202 sampling_batch_size=32768,203 stage="val",204 generator=self.test_generator,205 cfg=0,206 )207 samples = samples.view(batch_size * num_sample_per_cond, -1)208 batch_swarm = {"gps": samples, "emb": cond}209 nll_batch = -self.compute_exact_loglikelihood(batch_swarm, cfg=0)210 nll_batch = nll_batch.view(batch_size, num_sample_per_cond, -1)211 nll_best = nll_batch[212 torch.arange(batch_size), nll_batch.argmin(dim=1).squeeze(1)213 ]214 self.log(215 "test/best_nll",216 nll_best.mean(),217 sync_dist=True,218 on_step=False,219 on_epoch=True,220 )221 samples = samples.view(batch_size, num_sample_per_cond, -1)[222 torch.arange(batch_size), nll_batch.argmin(dim=1).squeeze(1)223 ]224 self.test_metrics.update({"gps": samples}, batch)225 226 def test_step(self, batch, batch_idx):227 if self.cfg.compute_swarms:228 self.test_best_nll(batch, batch_idx)229 else:230 self.test_step_simple(batch, batch_idx)231 232 def on_test_epoch_end(self):233 metrics = self.test_metrics.compute()234 for metric_name, metric_value in metrics.items():235 self.log(236 f"test/{metric_name}",237 metric_value,238 sync_dist=True,239 on_step=False,240 on_epoch=True,241 )242 243 def configure_optimizers(self):244 if self.cfg.optimizer.exclude_ln_and_biases_from_weight_decay:245 parameters_names_wd = get_parameter_names(self.network, [nn.LayerNorm])246 parameters_names_wd = [247 name for name in parameters_names_wd if "bias" not in name248 ]249 optimizer_grouped_parameters = [250 {251 "params": [252 p253 for n, p in self.network.named_parameters()254 if n in parameters_names_wd255 ],256 "weight_decay": self.cfg.optimizer.optim.weight_decay,257 "layer_adaptation": True,258 },259 {260 "params": [261 p262 for n, p in self.network.named_parameters()263 if n not in parameters_names_wd264 ],265 "weight_decay": 0.0,266 "layer_adaptation": False,267 },268 ]269 optimizer = instantiate(270 self.cfg.optimizer.optim, optimizer_grouped_parameters271 )272 else:273 optimizer = instantiate(self.cfg.optimizer.optim, self.network.parameters())274 if "lr_scheduler" in self.cfg:275 scheduler = instantiate(self.cfg.lr_scheduler)(optimizer)276 return [optimizer], [{"scheduler": scheduler, "interval": "step"}]277 else:278 return optimizer279 280 def lr_scheduler_step(self, scheduler, metric):281 scheduler.step(self.global_step)282 283 def sample(284 self,285 batch_size=None,286 cond=None,287 x_N=None,288 num_steps=None,289 stage="test",290 cfg=0,291 generator=None,292 return_trajectories=False,293 postprocessing=True,294 ):295 if x_N is None:296 assert batch_size is not None297 if isinstance(self.manifold, Sphere):298 x_N = self.manifold.random_base(299 batch_size, self.input_dim, device=self.device300 )301 x_N = x_N.reshape(batch_size, self.input_dim)302 else:303 x_N = torch.randn(batch_size, self.input_dim, device=self.device)304 batch = {"y": x_N}305 if stage == "val":306 sampler = self.val_sampler307 elif stage == "test":308 sampler = self.test_sampler309 else:310 raise ValueError(f"Unknown stage {stage}")311 batch[self.cfg.cond_preprocessing.input_key] = cond312 batch = self.cond_preprocessing(batch, device=self.device)313 if num_steps is None:314 output = sampler(315 self.ema_model,316 batch,317 conditioning_keys=self.cfg.cond_preprocessing.output_key,318 scheduler=self.inference_noise_scheduler,319 cfg_rate=cfg,320 generator=generator,321 return_trajectories=return_trajectories,322 )323 else:324 output = sampler(325 self.ema_model,326 batch,327 conditioning_keys=self.cfg.cond_preprocessing.output_key,328 scheduler=self.inference_noise_scheduler,329 num_steps=num_steps,330 cfg_rate=cfg,331 generator=generator,332 return_trajectories=return_trajectories,333 )334 if return_trajectories:335 return (336 self.postprocessing(output[0]) if postprocessing else output[0],337 [338 self.postprocessing(frame) if postprocessing else frame339 for frame in output[1]340 ],341 )342 else:343 return self.postprocessing(output) if postprocessing else output344 345 def sample_distribution(346 self,347 x_N,348 cond,349 sampling_batch_size=2048,350 num_steps=None,351 stage="test",352 cfg=0,353 generator=None,354 return_trajectories=False,355 ):356 if return_trajectories:357 x_0 = []358 trajectories = []359 i = -1360 for i in range(x_N.shape[0] // sampling_batch_size):361 x_N_batch = x_N[i * sampling_batch_size : (i + 1) * sampling_batch_size]362 cond_batch = cond[363 i * sampling_batch_size : (i + 1) * sampling_batch_size364 ]365 out, trajectories = self.sample(366 cond=cond_batch,367 x_N=x_N_batch,368 num_steps=num_steps,369 stage=stage,370 cfg=cfg,371 generator=generator,372 return_trajectories=return_trajectories,373 )374 x_0.append(out)375 trajectories.append(trajectories)376 if x_N.shape[0] % sampling_batch_size != 0:377 x_N_batch = x_N[(i + 1) * sampling_batch_size :]378 cond_batch = cond[(i + 1) * sampling_batch_size :]379 out, trajectories = self.sample(380 cond=cond_batch,381 x_N=x_N_batch,382 num_steps=num_steps,383 stage=stage,384 cfg=cfg,385 generator=generator,386 return_trajectories=return_trajectories,387 )388 x_0.append(out)389 trajectories.append(trajectories)390 x_0 = torch.cat(x_0, dim=1)391 trajectories = [torch.cat(frame, dim=1) for frame in trajectories]392 return x_0, trajectories393 else:394 x_0 = []395 i = -1396 for i in range(x_N.shape[0] // sampling_batch_size):397 x_N_batch = x_N[i * sampling_batch_size : (i + 1) * sampling_batch_size]398 cond_batch = cond[399 i * sampling_batch_size : (i + 1) * sampling_batch_size400 ]401 out = self.sample(402 cond=cond_batch,403 x_N=x_N_batch,404 num_steps=num_steps,405 stage=stage,406 cfg=cfg,407 generator=generator,408 return_trajectories=return_trajectories,409 )410 x_0.append(out)411 if x_N.shape[0] % sampling_batch_size != 0:412 x_N_batch = x_N[(i + 1) * sampling_batch_size :]413 cond_batch = cond[(i + 1) * sampling_batch_size :]414 out = self.sample(415 cond=cond_batch,416 x_N=x_N_batch,417 num_steps=num_steps,418 stage=stage,419 cfg=cfg,420 generator=generator,421 return_trajectories=return_trajectories,422 )423 x_0.append(out)424 x_0 = torch.cat(x_0, dim=0)425 return x_0426 427 def model(self, *args, **kwargs):428 return self.preconditioning(self.network, *args, **kwargs)429 430 def ema_model(self, *args, **kwargs):431 return self.preconditioning(self.ema_network, *args, **kwargs)432 433 def compute_exact_loglikelihood(434 self,435 batch=None,436 x_1=None,437 cond=None,438 t1=1.0,439 num_steps=1000,440 rademacher=False,441 data_preprocessing=True,442 cfg=0,443 ):444 nfe = [0]445 if batch is None:446 batch = {"x_0": x_1, "emb": cond}447 if data_preprocessing:448 batch = self.data_preprocessing(batch)449 batch = self.cond_preprocessing(batch)450 timesteps = self.inference_noise_scheduler(451 torch.linspace(0, t1, 2).to(batch["x_0"])452 )453 with torch.inference_mode(mode=False):454 455 def odefunc(t, tensor):456 nfe[0] += 1457 t = t.to(tensor)458 gamma = self.inference_noise_scheduler(t)459 x = tensor[..., : self.input_dim]460 y = batch["emb"]461 462 def vecfield(x, y):463 if cfg > 0:464 batch_vecfield = {465 "y": x,466 "emb": y,467 "gamma": gamma.reshape(-1),468 }469 model_output_cond = self.ema_model(batch_vecfield)470 batch_vecfield_uncond = {471 "y": x,472 "emb": torch.zeros_like(y),473 "gamma": gamma.reshape(-1),474 }475 model_output_uncond = self.ema_model(batch_vecfield_uncond)476 model_output = model_output_cond + cfg * (477 model_output_cond - model_output_uncond478 )479 480 else:481 batch_vecfield = {482 "y": x,483 "emb": y,484 "gamma": gamma.reshape(-1),485 }486 model_output = self.ema_model(batch_vecfield)487 488 if self.interpolant == "flow_matching":489 d_gamma = self.inference_noise_scheduler.derivative(t).reshape(490 -1, 1491 )492 return d_gamma * model_output493 elif self.interpolant == "diffusion":494 alpha_t = self.inference_noise_scheduler.alpha(t).reshape(-1, 1)495 return (496 -1 / 2 * (alpha_t * x - torch.abs(alpha_t) * model_output)497 )498 else:499 raise ValueError(f"Unknown interpolant {self.interpolant}")500 501 if rademacher:502 v = torch.randint_like(x, 2) * 2 - 1503 else:504 v = None505 dx, div = output_and_div(vecfield, x, y, v=v)506 div = div.reshape(-1, 1)507 del t, x508 return torch.cat([dx, div], dim=-1)509 510 x_1 = batch["x_0"]511 state1 = torch.cat([x_1, torch.zeros_like(x_1[..., :1])], dim=-1)512 with torch.no_grad():513 if False and isinstance(self.manifold, Sphere):514 print("Riemannian flow sampler")515 product_man = ProductManifold(516 (self.manifold, self.input_dim), (Euclidean(), 1)517 )518 state0 = ode_riemannian_flow_sampler(519 odefunc,520 state1,521 manifold=product_man,522 scheduler=self.inference_noise_scheduler,523 num_steps=num_steps,524 )525 else:526 print("ODE solver")527 state0 = odeint(528 odefunc,529 state1,530 t=torch.linspace(0, t1, 2).to(batch["x_0"]),531 atol=1e-6,532 rtol=1e-6,533 method="dopri5",534 options={"min_step": 1e-5},535 )[-1]536 x_0, logdetjac = state0[..., : self.input_dim], state0[..., -1]537 if self.manifold is not None:538 x_0 = self.manifold.projx(x_0)539 logp0 = self.manifold.base_logprob(x_0)540 else:541 logp0 = (542 -1 / 2 * (x_0**2).sum(dim=-1)543 - self.input_dim544 * torch.log(torch.tensor(2 * np.pi, device=x_0.device))545 / 2546 )547 print(f"nfe: {nfe[0]}")548 logp1 = logp0 + logdetjac549 logp1 = logp1 / (self.input_dim * np.log(2))550 return logp1551 552 553def get_parameter_names(model, forbidden_layer_types):554 """555 Returns the names of the model parameters that are not inside a forbidden layer.556 Taken from HuggingFace transformers.557 """558 result = []559 for name, child in model.named_children():560 result += [561 f"{name}.{n}"562 for n in get_parameter_names(child, forbidden_layer_types)563 if not isinstance(child, tuple(forbidden_layer_types))564 ]565 # Add model specific parameters (defined with nn.Parameter) since they are not in any child.566 result += list(model._parameters.keys())567 return result568 569 570# for likelihood computation571def div_fn(u):572 """Accepts a function u:R^D -> R^D."""573 J = jacrev(u, argnums=0)574 return lambda x, y: torch.trace(J(x, y).squeeze(0))575 576 577def output_and_div(vecfield, x, y, v=None):578 if v is None:579 dx = vecfield(x, y)580 div = vmap(div_fn(vecfield))(x, y)581 else:582 vecfield_x = lambda x: vecfield(x, y)583 dx, vjpfunc = vjp(vecfield_x, x)584 vJ = vjpfunc(v)[0]585 div = torch.sum(vJ * v, dim=-1)586 return dx, div587 588 589class VonFisherGeolocalizer(L.LightningModule):590 def __init__(self, cfg):591 super().__init__()592 self.cfg = cfg593 self.network = instantiate(cfg.network)594 # self.network = torch.compile(self.network, fullgraph=True)595 self.input_dim = cfg.network.input_dim596 self.data_preprocessing = instantiate(cfg.data_preprocessing)597 self.cond_preprocessing = instantiate(cfg.cond_preprocessing)598 self.preconditioning = instantiate(cfg.preconditioning)599 600 self.ema_network = copy.deepcopy(self.network).requires_grad_(False)601 self.ema_network.eval()602 self.postprocessing = instantiate(cfg.postprocessing)603 self.val_sampler = instantiate(cfg.val_sampler)604 self.test_sampler = instantiate(cfg.test_sampler)605 self.loss = instantiate(cfg.loss)()606 self.val_metrics = instantiate(cfg.val_metrics)607 self.test_metrics = instantiate(cfg.test_metrics)608 609 def training_step(self, batch, batch_idx):610 with torch.no_grad():611 batch = self.data_preprocessing(batch)612 batch = self.cond_preprocessing(batch)613 batch_size = batch["x_0"].shape[0]614 loss = self.loss(self.preconditioning, self.network, batch).mean()615 self.log(616 "train/loss",617 loss,618 sync_dist=True,619 on_step=True,620 on_epoch=True,621 batch_size=batch_size,622 )623 return loss624 625 def on_before_optimizer_step(self, optimizer):626 if self.global_step == 0:627 no_grad = []628 for name, param in self.network.named_parameters():629 if param.grad is None:630 no_grad.append(name)631 if len(no_grad) > 0:632 print("Parameters without grad:")633 print(no_grad)634 635 def on_validation_start(self):636 self.validation_generator = torch.Generator(device=self.device).manual_seed(637 3407638 )639 self.validation_generator_ema = torch.Generator(device=self.device).manual_seed(640 3407641 )642 643 def validation_step(self, batch, batch_idx):644 batch = self.data_preprocessing(batch)645 batch = self.cond_preprocessing(batch)646 batch_size = batch["x_0"].shape[0]647 loss = self.loss(648 self.preconditioning,649 self.network,650 batch,651 generator=self.validation_generator,652 ).mean()653 self.log(654 "val/loss",655 loss,656 sync_dist=True,657 on_step=False,658 on_epoch=True,659 batch_size=batch_size,660 )661 if hasattr(self, "ema_model"):662 loss_ema = self.loss(663 self.preconditioning,664 self.ema_network,665 batch,666 generator=self.validation_generator_ema,667 ).mean()668 self.log(669 "val/loss_ema",670 loss_ema,671 sync_dist=True,672 on_step=False,673 on_epoch=True,674 batch_size=batch_size,675 )676 677 def on_test_start(self):678 self.test_generator = torch.Generator(device=self.device).manual_seed(3407)679 680 def test_step(self, batch, batch_idx):681 batch = self.data_preprocessing(batch)682 batch = self.cond_preprocessing(batch)683 batch_size = batch["x_0"].shape[0]684 cond = batch[self.cfg.cond_preprocessing.output_key]685 686 samples = self.sample(cond=cond, stage="test")687 self.test_metrics.update({"gps": samples}, batch)688 nll = -self.compute_exact_loglikelihood(batch).mean()689 self.log(690 "test/NLL",691 nll,692 sync_dist=True,693 on_step=False,694 on_epoch=True,695 batch_size=batch_size,696 )697 698 def on_test_epoch_end(self):699 metrics = self.test_metrics.compute()700 for metric_name, metric_value in metrics.items():701 self.log(702 f"test/{metric_name}",703 metric_value,704 sync_dist=True,705 on_step=False,706 on_epoch=True,707 )708 709 def configure_optimizers(self):710 if self.cfg.optimizer.exclude_ln_and_biases_from_weight_decay:711 parameters_names_wd = get_parameter_names(self.network, [nn.LayerNorm])712 parameters_names_wd = [713 name for name in parameters_names_wd if "bias" not in name714 ]715 optimizer_grouped_parameters = [716 {717 "params": [718 p719 for n, p in self.network.named_parameters()720 if n in parameters_names_wd721 ],722 "weight_decay": self.cfg.optimizer.optim.weight_decay,723 "layer_adaptation": True,724 },725 {726 "params": [727 p728 for n, p in self.network.named_parameters()729 if n not in parameters_names_wd730 ],731 "weight_decay": 0.0,732 "layer_adaptation": False,733 },734 ]735 optimizer = instantiate(736 self.cfg.optimizer.optim, optimizer_grouped_parameters737 )738 else:739 optimizer = instantiate(self.cfg.optimizer.optim, self.network.parameters())740 if "lr_scheduler" in self.cfg:741 scheduler = instantiate(self.cfg.lr_scheduler)(optimizer)742 return [optimizer], [{"scheduler": scheduler, "interval": "step"}]743 else:744 return optimizer745 746 def lr_scheduler_step(self, scheduler, metric):747 scheduler.step(self.global_step)748 749 def sample(750 self,751 batch_size=None,752 cond=None,753 postprocessing=True,754 stage="val",755 ):756 batch = {}757 if stage == "val":758 sampler = self.val_sampler759 elif stage == "test":760 sampler = self.test_sampler761 else:762 raise ValueError(f"Unknown stage {stage}")763 batch[self.cfg.cond_preprocessing.input_key] = cond764 batch = self.cond_preprocessing(batch, device=self.device)765 output = sampler(766 self.ema_model,767 batch,768 )769 return self.postprocessing(output) if postprocessing else output770 771 def model(self, *args, **kwargs):772 return self.preconditioning(self.network, *args, **kwargs)773 774 def ema_model(self, *args, **kwargs):775 return self.preconditioning(self.ema_network, *args, **kwargs)776 777 def compute_exact_loglikelihood(778 self,779 batch=None,780 ):781 batch = self.data_preprocessing(batch)782 batch = self.cond_preprocessing(batch)783 return -self.loss(self.preconditioning, self.ema_network, batch)784 785 786class RandomGeolocalizer(L.LightningModule):787 def __init__(self, cfg):788 super().__init__()789 self.cfg = cfg790 self.test_metrics = instantiate(cfg.test_metrics)791 self.data_preprocessing = instantiate(cfg.data_preprocessing)792 self.cond_preprocessing = instantiate(cfg.cond_preprocessing)793 self.postprocessing = instantiate(cfg.postprocessing)794 795 def test_step(self, batch, batch_idx):796 batch = self.data_preprocessing(batch)797 batch = self.cond_preprocessing(batch)798 batch_size = batch["x_0"].shape[0]799 samples = torch.randn(batch_size, 3, device=self.device)800 samples = samples / samples.norm(dim=-1, keepdim=True)801 samples = self.postprocessing(samples)802 self.test_metrics.update({"gps": samples}, batch)803 804 def on_test_epoch_end(self):805 metrics = self.test_metrics.compute()806 for metric_name, metric_value in metrics.items():807 self.log(808 f"test/{metric_name}",809 metric_value,810 sync_dist=True,811 on_step=False,812 on_epoch=True,813 )814