def ssm_loss_fn()

in jcm/losses.py [0:0]

• 77 lines of code
• 4 McCabe index (conditional complexity)

    def ssm_loss_fn(rng, params, states, batch):
        """Compute the loss function based on sliced score matching.

        Args:
          rng: A JAX random state.
          params: A dictionary that contains trainable parameters of the score-based model.
          states: A dictionary that contains mutable states of the score-based model.
          batch: A mini-batch of training data.

        Returns:
          loss: A scalar that represents the average loss value across the mini-batch.
          new_model_state: A dictionary that contains the mutated states of the score-based model.
        """

        score_fn = mutils.get_score_fn(
            sde,
            model,
            params,
            states,
            train=train,
            return_state=True,
        )
        data = batch["image"]
        rng = hk.PRNGSequence(rng)
        # DEBUG: beware of eps!
        if isinstance(sde, sde_lib.KVESDE):
            t = random.normal(next(rng), (data.shape[0],)) * 1.2 - 1.2
            t = jnp.exp(t)
        else:
            t = random.uniform(next(rng), (data.shape[0],), minval=eps, maxval=sde.T)
        # t = random.uniform(next(rng), (data.shape[0],), minval=eps, maxval=sde.T)
        z = random.normal(next(rng), data.shape)
        mean, std = sde.marginal_prob(data, t)
        perturbed_data = mean + batch_mul(std, z)

        def score_fn_for_jvp(x):
            return score_fn(x, t, rng=next(rng))

        epsilon = random.rademacher(next(rng), data.shape, dtype=data.dtype)
        score, score_trace, new_model_state = jax.jvp(
            score_fn_for_jvp, (perturbed_data,), (epsilon,), has_aux=True
        )
        score_norm = jnp.mean(jnp.square(score).reshape((score.shape[0], -1)), axis=-1)
        score_trace = jnp.mean(
            (2 * score_trace * epsilon).reshape((score.shape[0], -1)), axis=-1
        )

        if not likelihood_weighting:
            losses = (score_norm + score_trace) * std**2
        elif isinstance(sde, sde_lib.KVESDE):
            losses = score_norm + score_trace
            losses = (
                losses * std**2 * (std**2 + sde.data_std**2) / sde.data_std**2
            )
        else:
            g2 = sde.sde(jnp.zeros_like(data), t)[1] ** 2
            losses = (score_norm + score_trace) * g2
        loss = jnp.nansum(losses * batch["mask"] / jnp.sum(batch["mask"]))
        quarter_masks = get_quarter_masks(
            t,
            np.linspace(sde.t_min ** (1 / sde.rho), sde.t_max ** (1 / sde.rho), 5)
            ** sde.rho,
        )
        loss_q1 = jnp.nansum(
            losses
            * quarter_masks[0]
            * batch["mask"]
            / jnp.sum(quarter_masks[0] * batch["mask"])
        )
        loss_q2 = jnp.nansum(
            losses
            * quarter_masks[1]
            * batch["mask"]
            / jnp.sum(quarter_masks[1] * batch["mask"])
        )
        loss_q3 = jnp.nansum(
            losses
            * quarter_masks[2]
            * batch["mask"]
            / jnp.sum(quarter_masks[2] * batch["mask"])
        )
        loss_q4 = jnp.nansum(
            losses
            * quarter_masks[3]
            * batch["mask"]
            / jnp.sum(quarter_masks[3] * batch["mask"])
        )

        log_stats = {
            "loss_q1": loss_q1,
            "loss_q2": loss_q2,
            "loss_q3": loss_q3,
            "loss_q4": loss_q4,
            "loss": loss,
        }

        return loss, (new_model_state, log_stats)