def loss_fn()

in jcm/losses.py [0:0]

• 73 lines of code
• 9 McCabe index (conditional complexity)

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

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

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

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

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

        dropout_rng = next(rng)
        denoised_x, new_states = mutils.get_denoiser_fn(
            sde, model, params, states, train=train, return_state=True
        )(x_t, t, rng=dropout_rng if train else None)

        target_denoiser_fn = mutils.get_denoiser_fn(
            sde,
            model,
            target_params,
            states,
            train=False,
            return_state=False,
        )

        def euler_solver(samples, t, next_t):
            x = samples
            denoiser = target_denoiser_fn(x, t, rng=None)
            score = batch_mul(1 / t**2, denoiser - x)
            samples = x + batch_mul(next_t - t, -batch_mul(score, t))

            return samples

        def euler_to_denoiser(x_t, t, x_next_t, next_t):
            denoiser = x_t - batch_mul(t, batch_mul(x_next_t - x_t, 1 / (next_t - t)))
            return denoiser

        x_t2 = euler_solver(x_t, t, t2)
        x_t3 = euler_solver(x_t2, t2, t3)

        target_x = jax.lax.stop_gradient(euler_to_denoiser(x_t, t, x_t3, t3))

        diffs = denoised_x - target_x

        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_denoised_x = jax.image.resize(
                denoised_x, (denoised_x.shape[0], 224, 224, 3), method="bilinear"
            )
            scaled_target_x = jax.image.resize(
                target_x, (target_x.shape[0], 224, 224, 3), method="bilinear"
            )
            losses = jnp.squeeze(
                lpips_model.apply(lpips_params, scaled_denoised_x, scaled_target_x)
            )
        else:
            raise ValueError("Unknown loss norm: {}".format(loss_norm))

        if weighting.lower() == "snrp1":
            weight = 1 / t**2 + 1
        elif weighting.lower() == "truncated_snr":
            weight = jnp.maximum(1 / t**2, jnp.ones_like(t))
        elif weighting.lower() == "snr":
            weight = 1 / t**2

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

        return loss, (new_states, log_stats)