def loss_fn()

in jcm/losses.py [0:0]

• 127 lines of code
• 10 McCabe index (conditional complexity)

    def loss_fn(rng, params, states, batch):
        rng = hk.PRNGSequence(rng)
        x = batch["image"]

        # sampling t according to the Heun sampler
        t = jax.random.uniform(
            next(rng),
            (x.shape[0],),
            minval=sde.t_min ** (1 / sde.rho),
            maxval=sde.t_max ** (1 / sde.rho),
        ) ** (sde.rho)

        weightings = jnp.ones_like(t)
        z = jax.random.normal(next(rng), x.shape)
        x_t = x + batch_mul(t, z)

        if dsm_target:
            score_t = batch_mul(x - x_t, 1 / t**2)
        else:
            score_t = score_fn(x_t, t)

        if train:
            step_rng = next(rng)
        else:
            step_rng = None

        def model_fn(data, time):
            return mutils.get_distiller_fn(
                sde, model, params, states, train=train, return_state=True
            )(data, time, rng=step_rng)

        Ft, diffs, new_states = jax.jvp(
            model_fn, (x_t, t), (batch_mul(t, score_t), -jnp.ones_like(t)), has_aux=True
        )

        if loss_norm.lower() == "l1":
            losses = jnp.abs(diffs)
            losses = jnp.mean(losses.reshape(losses.shape[0], -1), axis=1)
        elif loss_norm.lower() == "l2":
            losses = diffs**2
            losses = jnp.sqrt(jnp.sum(losses.reshape(losses.shape[0], -1), axis=1))
        elif loss_norm.lower() == "linf":
            losses = jnp.abs(diffs)
            losses = jnp.max(losses.reshape(losses.shape[0], -1), axis=1)
        elif loss_norm.lower() == "lpips":

            def metric(x):
                scaled_Ft = jax.image.resize(
                    Ft, (Ft.shape[0], 224, 224, 3), method="bilinear"
                )
                x = jax.image.resize(x, (x.shape[0], 224, 224, 3), method="bilinear")
                return jnp.sum(
                    jnp.squeeze(lpips_model.apply(lpips_params, scaled_Ft, x))
                )

            losses = (
                jax.grad(lambda x: jnp.sum(jax.grad(metric)(x) * diffs))(Ft) * diffs
            )
            losses = jnp.sum(losses.reshape(losses.shape[0], -1), axis=1)

        else:
            raise ValueError("Unknown loss norm: {}".format(loss_norm))

        if stopgrad:
            if loss_norm.lower() == "l2":
                pseudo_losses = -jax.lax.stop_gradient(diffs) * Ft
                pseudo_losses = jnp.sum(
                    pseudo_losses.reshape((pseudo_losses.shape[0], -1)), axis=-1
                )
                loss = jnp.nansum(
                    pseudo_losses * batch["mask"] * weightings / jnp.sum(batch["mask"])
                )
            elif loss_norm.lower() == "lpips":

                def metric_fn(x):
                    x = jax.image.resize(
                        x, (x.shape[0], 224, 224, 3), method="bilinear"
                    )
                    y = jax.image.resize(
                        jax.lax.stop_gradient(Ft),
                        (x.shape[0], 224, 224, 3),
                        method="bilinear",
                    )
                    return jnp.sum(jnp.squeeze(lpips_model.apply(lpips_params, x, y)))

                # forward-over-reverse
                def hvp(f, primals, tangents):
                    return jax.jvp(jax.grad(f), primals, tangents)[1]

                pseudo_losses = Ft * hvp(
                    metric_fn,
                    (jax.lax.stop_gradient(Ft),),
                    (-jax.lax.stop_gradient(diffs),),
                )
                pseudo_losses = jnp.sum(
                    pseudo_losses.reshape((pseudo_losses.shape[0], -1)), axis=-1
                )
                loss = jnp.nansum(
                    pseudo_losses * batch["mask"] * weightings / jnp.sum(batch["mask"])
                )
            else:
                raise NotImplementedError

        else:
            loss = jnp.nansum(
                losses * batch["mask"] * weightings / 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": loss,
            "loss_q1": loss_q1,
            "loss_q2": loss_q2,
            "loss_q3": loss_q3,
            "loss_q4": loss_q4,
        }

        return loss, (new_states, log_stats)