Method training_losses

diff2flow/flow_obj.py:192–216

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. cond_kwargs: additional arguments for the conditional flow

(self, x1: Tensor, x0: Tensor = None, **cond_kwargs)

    def training_losses(self, x1: Tensor, x0: Tensor = None, **cond_kwargs):
        """
        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.
            cond_kwargs: additional arguments for the conditional flow
                network (e.g. conditioning information)
        Returns:
            loss: scalar, the training loss for the flow model
        """
        if x0 is None:
            x0 = torch.randn_like(x1)

        bs, dev, dtype = x1.shape[0], x1.device, x1.dtype

        # Sample time t from uniform distribution U(0, 1)
        t = torch.rand(bs, device=dev, dtype=dtype)

        # sample xt and ut
        xt = self.compute_xt(x0=x0, x1=x1, t=t)
        ut = self.compute_ut(x0=x0, x1=x1, t=t)
        vt = self.sample_vt(fm_x=xt, fm_t=t, **cond_kwargs)

        return (vt - ut).square()