def loss_fn()

in jcm/losses.py [0:0]

• 61 lines of code
• 15 McCabe index (conditional complexity)

    def loss_fn(rng, params, states, batch, target_params=None, num_scales=None):
        rng = hk.PRNGSequence(rng)
        x = batch["image"]
        if target_params is None:
            target_params = params

        if num_scales is None:
            num_scales = sde.N

        indices = jax.random.randint(next(rng), (x.shape[0],), 0, num_scales - 1)
        t = sde.t_max ** (1 / sde.rho) + indices / (num_scales - 1) * (
            sde.t_min ** (1 / sde.rho) - sde.t_max ** (1 / sde.rho)
        )
        t = t**sde.rho

        t2 = sde.t_max ** (1 / sde.rho) + (indices + 1) / (num_scales - 1) * (
            sde.t_min ** (1 / sde.rho) - sde.t_max ** (1 / sde.rho)
        )
        t2 = t2**sde.rho

        z = jax.random.normal(next(rng), x.shape)
        x_t = x + batch_mul(t, z)
        dropout_rng = next(rng)
        Ft, new_states = mutils.get_distiller_fn(
            sde, model, params, states, train=train, return_state=True
        )(x_t, t, rng=dropout_rng if train else None)

        x_t2 = ode_solver(x_t, t, t2, x)
        Ft2, new_states = mutils.get_distiller_fn(
            sde, model, target_params, new_states, train=train, return_state=True
        )(x_t2, t2, rng=dropout_rng if train else None)

        if stopgrad:
            Ft2 = jax.lax.stop_gradient(Ft2)

        diffs = Ft - Ft2

        if weighting.lower() == "uniform":
            weight = jnp.ones_like(t)
        elif weighting.lower() == "snrp1":
            weight = 1 / t**2 + 1.0
        elif weighting.lower() == "truncated_snr":
            weight = jnp.maximum(1 / t**2, jnp.ones_like(t))
        elif weighting.lower() == "snr":
            weight = 1 / t**2
        else:
            raise NotImplementedError(f"Weighting {weighting} not implemented")

        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.mean(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":
            scaled_Ft = jax.image.resize(
                Ft, (Ft.shape[0], 224, 224, 3), method="bilinear"
            )
            scaled_Ft2 = jax.image.resize(
                Ft2, (Ft2.shape[0], 224, 224, 3), method="bilinear"
            )
            losses = jnp.squeeze(lpips_model.apply(lpips_params, scaled_Ft, scaled_Ft2))

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

        loss = jnp.nansum(losses * batch["mask"] * weight / jnp.sum(batch["mask"]))
        log_stats = {}

        ## Uncomment to log loss per time step
        # for t_index in range(sde.N - 1):
        #     mask = (indices == t_index).astype(jnp.float32)
        #     log_stats["loss_t{}".format(t_index)] = jnp.nansum(
        #         losses * batch["mask"] * mask / jnp.sum(batch["mask"] * mask)
        #     )

        return loss, (new_states, log_stats)