Args: x1: shape (bs, *dim), represents the target minibatch (data) x0: shape (bs, *dim), represents the source minibatch, if None we sample x0 from a standard normal distribution. t: shape (bs,), represents the time in [0, 1]. If None, we
(self, x1: Tensor, x0: Tensor = None, **cond_kwargs)
| 589 | return alpha_dt_t * x1 + sigma_dt_t * x0 |
| 590 | |
| 591 | def training_losses(self, x1: Tensor, x0: Tensor = None, **cond_kwargs): |
| 592 | """ |
| 593 | Args: |
| 594 | x1: shape (bs, *dim), represents the target minibatch (data) |
| 595 | x0: shape (bs, *dim), represents the source minibatch, if None |
| 596 | we sample x0 from a standard normal distribution. |
| 597 | t: shape (bs,), represents the time in [0, 1]. If None, we sample |
| 598 | according to the t_sampler (default: U(0, 1)). |
| 599 | cond_kwargs: additional arguments for the conditional flow |
| 600 | network (e.g. conditioning information) |
| 601 | Returns: |
| 602 | loss: scalar, the training loss for the flow model |
| 603 | """ |
| 604 | if x0 is None: x0 = torch.randn_like(x1) |
| 605 | t = self.t_sampler(x1.shape[0], device=x1.device, dtype=x1.dtype) |
| 606 | |
| 607 | xt = self.compute_xt(x0=x0, x1=x1, t=t) |
| 608 | ut = self.compute_ut(x0=x0, x1=x1, t=t) |
| 609 | vt = self.forward(x=xt, t=t, **cond_kwargs) |
| 610 | |
| 611 | return (vt - ut).square() |
| 612 | |
| 613 | def validation_losses(self, x1: Tensor, x0: Tensor = None, num_segments: int = 8, **cond_kwargs): |
| 614 | """ |
no test coverage detected