def dsm_loss_fn()

in jcm/losses.py [0:0]

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

    def dsm_loss_fn(rng, params, states, batch):
        """Compute the loss function based on denoising 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.
        """

        data = batch["image"]
        rng = hk.PRNGSequence(rng)

        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)

        z = random.normal(next(rng), data.shape)
        mean, std = sde.marginal_prob(data, t)
        perturbed_data = mean + batch_mul(std, z)

        if isinstance(sde, sde_lib.KVESDE):
            score_fn = mutils.get_score_fn(
                sde,
                model,
                params,
                states,
                train=train,
                return_state=True,
            )
            score, new_model_state = score_fn(perturbed_data, t, rng=next(rng))

            losses = jnp.square(batch_mul(score, std) + z)
            losses = batch_mul(
                losses, (std**2 + sde.data_std**2) / sde.data_std**2
            )
            losses = jnp.sum(losses.reshape((losses.shape[0], -1)), axis=-1)

        else:
            score_fn = mutils.get_score_fn(
                sde,
                model,
                params,
                states,
                train=train,
                return_state=True,
            )

            score, new_model_state = score_fn(perturbed_data, t, rng=next(rng))

            if not likelihood_weighting:
                losses = jnp.square(batch_mul(score, std) + z)
                losses = jnp.mean(losses.reshape((losses.shape[0], -1)), axis=-1)

            else:
                g2 = sde.sde(jnp.zeros_like(data), t)[1] ** 2
                losses = jnp.square(score + batch_mul(z, 1.0 / std))
                losses = jnp.mean(losses.reshape((losses.shape[0], -1)), axis=-1) * 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,
        }

        return loss, (new_model_state, log_stats)