def get_loss_fn()

in jcm/losses.py [0:0]

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

def get_loss_fn(config, sde, score_model, state, rng):
    likelihood_weighting = config.training.likelihood_weighting
    if config.training.loss.lower() in ["dsm", "ssm"]:
        ssm = config.training.loss.lower() == "ssm"
        train_loss_fn = get_score_matching_loss_fn(
            sde,
            score_model,
            train=True,
            likelihood_weighting=likelihood_weighting,
            ssm=ssm,
        )
        eval_loss_fn = get_score_matching_loss_fn(
            sde,
            score_model,
            train=False,
            likelihood_weighting=likelihood_weighting,
            ssm=ssm,
        )
    elif config.training.loss.lower().startswith(
        ("continuous", "consistency", "progressive_distillation")
    ):
        optimizer, optimize_fn = get_optimizer(config.training.ref_config)
        rng = hk.PRNGSequence(rng)
        ref_config = config.training.ref_config
        ref_model, init_ref_model_state, init_ref_params = mutils.init_model(
            next(rng), ref_config
        )
        ref_state = mutils.State(
            step=0,
            lr=ref_config.optim.lr,
            ema_rate=ref_config.model.ema_rate,
            params=init_ref_params,
            params_ema=init_ref_params,
            model_state=init_ref_model_state,
            opt_state=optimizer.init(init_ref_params),
            rng_state=rng.internal_state,
        )
        ref_state = checkpoints.restore_checkpoint(
            config.training.ref_model_path, ref_state
        )
        # Initialize the flow model from the denoiser model
        if config.training.finetune:
            state = state.replace(
                params=ref_state.params,
                params_ema=ref_state.params_ema,
                model_state=ref_state.model_state,
            )
        if config.training.loss_norm.lower() == "lpips":
            lpips_model, lpips_params = mutils.init_lpips(next(rng), config)
        else:
            lpips_model, lpips_params = None, None
        if config.training.loss.lower().startswith("continuous"):
            train_loss_fn = get_continuous_consistency_loss_fn(
                sde,
                ref_model,
                ref_state.params_ema,
                ref_state.model_state,
                score_model,
                train=True,
                loss_norm=config.training.loss_norm,
                stopgrad=config.training.stopgrad,
                lpips_model=lpips_model,
                lpips_params=lpips_params,
                dsm_target=config.training.dsm_target,
            )
            eval_loss_fn = get_continuous_consistency_loss_fn(
                sde,
                ref_model,
                ref_state.params_ema,
                ref_state.model_state,
                score_model,
                train=False,
                loss_norm=config.training.loss_norm,
                stopgrad=config.training.stopgrad,
                lpips_model=lpips_model,
                lpips_params=lpips_params,
                dsm_target=config.training.dsm_target,
            )
        elif config.training.loss.lower().startswith("consistency"):
            train_loss_fn = get_consistency_loss_fn(
                sde,
                ref_model,
                ref_state.params_ema,
                ref_state.model_state,
                score_model,
                train=True,
                loss_norm=config.training.loss_norm,
                weighting=config.training.weighting,
                stopgrad=config.training.stopgrad,
                dsm_target=config.training.dsm_target,
                solver=config.training.solver,
                lpips_model=lpips_model,
                lpips_params=lpips_params,
            )
            eval_loss_fn = get_consistency_loss_fn(
                sde,
                ref_model,
                ref_state.params_ema,
                ref_state.model_state,
                score_model,
                train=False,
                loss_norm=config.training.loss_norm,
                weighting=config.training.weighting,
                stopgrad=config.training.stopgrad,
                dsm_target=config.training.dsm_target,
                solver=config.training.solver,
                lpips_model=lpips_model,
                lpips_params=lpips_params,
            )

        elif config.training.loss.lower() == "progressive_distillation":
            train_loss_fn = get_progressive_distillation_loss_fn(
                sde,
                score_model,
                train=True,
                loss_norm=config.training.loss_norm,
                weighting=config.training.weighting,
                lpips_model=lpips_model,
                lpips_params=lpips_params,
            )
            eval_loss_fn = get_progressive_distillation_loss_fn(
                sde,
                score_model,
                train=False,
                loss_norm=config.training.loss_norm,
                weighting=config.training.weighting,
                lpips_model=lpips_model,
                lpips_params=lpips_params,
            )
            assert (
                config.training.finetune
            ), "Finetuning is required for progressive distillation."
            state = state.replace(
                target_params=ref_state.params_ema,
            )

    else:
        raise ValueError(f"Unknown loss {config.training.loss}")

    return train_loss_fn, eval_loss_fn, state