def get_ema_scales_fn()

in jcm/losses.py [0:0]

• 54 lines of code
• 8 McCabe index (conditional complexity)

def get_ema_scales_fn(config):
    if config.training.loss.lower() in ("dsm", "ssm", "continuous", "consistency"):

        def ema_and_scales_fn(step):
            return None, None

    else:

        def ema_and_scales_fn(step):
            if (
                config.training.target_ema_mode == "fixed"
                and config.training.scale_mode == "fixed"
            ):
                target_ema = float(config.training.target_ema)
                scales = int(config.model.num_scales)

            elif (
                config.training.target_ema_mode == "adaptive"
                and config.training.scale_mode == "progressive"
            ):
                start_ema = float(config.training.start_ema)
                start_scales = int(config.training.start_scales)
                end_scales = int(config.training.end_scales)
                total_steps = int(config.training.n_iters)
                scales = jnp.ceil(
                    jnp.sqrt(
                        (step / total_steps)
                        * ((end_scales + 1) ** 2 - start_scales**2)
                        + start_scales**2
                    )
                    - 1
                ).astype(jnp.int32)
                scales = jnp.maximum(scales, 1)
                c = -jnp.log(start_ema) * start_scales
                target_ema = jnp.exp(-c / scales)
                scales = scales + 1
            elif (
                config.training.target_ema_mode == "fixed"
                and config.training.scale_mode == "progdist"
            ):
                start_scales = int(config.training.start_scales)
                distill_steps_per_iter = int(config.training.distill_steps_per_iter)
                distill_stage = step // distill_steps_per_iter
                scales = start_scales // (2**distill_stage)
                scales = jnp.maximum(scales, 2)

                sub_stage = jnp.maximum(
                    step - distill_steps_per_iter * (jnp.log2(start_scales) - 1),
                    0,
                )
                sub_stage = sub_stage // (distill_steps_per_iter * 2)
                sub_scales = 2 // (2**sub_stage)
                sub_scales = jnp.maximum(sub_scales, 1)

                scales = jnp.where(scales == 2, sub_scales, scales)

                target_ema = 1.0
            else:
                raise NotImplementedError

            return target_ema, scales

    return ema_and_scales_fn