jcm/evaluate.py

# Copyright 2023 (c) OpenAI.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import io
import os
import time
from typing import Any

import flax
import jax
import jax.numpy as jnp
import numpy as np
import logging
import functools
import haiku as hk
import math
from collections import defaultdict

from . import checkpoints

# Keep the import below for registering all model definitions
from .models import ddpm, ncsnv2, ncsnpp

from .models import utils as mutils
from . import losses
from . import sampling
from . import datasets
from . import metrics
from . import likelihood
from . import sde_lib
from .metrics import get_samples_from_ckpt
import blobfile


def evaluate(config, workdir, eval_folder="eval"):
    """Evaluate trained models.

    Args:
      config: Configuration to use.
      workdir: Working directory for checkpoints.
      eval_folder: The subfolder for storing evaluation results. Default to
        "eval".
    """
    # Create directory to eval_folder
    eval_dir = os.path.join(workdir, eval_folder)
    blobfile.makedirs(eval_dir)

    rng = hk.PRNGSequence(config.seed + 1)
    # Initialize model
    score_model, init_model_state, initial_params = mutils.init_model(next(rng), config)
    optimizer, optimize_fn = losses.get_optimizer(config)
    if config.training.loss.lower().endswith(
        ("ema", "adaptive", "progressive_distillation")
    ):
        state = mutils.StateWithTarget(
            step=0,
            lr=config.optim.lr,
            ema_rate=config.model.ema_rate,
            params=initial_params,
            target_params=initial_params,
            params_ema=initial_params,
            model_state=init_model_state,
            opt_state=optimizer.init(initial_params),
            rng_state=rng.internal_state,
        )
    else:
        state = mutils.State(
            step=0,
            lr=config.optim.lr,
            ema_rate=config.model.ema_rate,
            params=initial_params,
            params_ema=initial_params,
            model_state=init_model_state,
            opt_state=optimizer.init(initial_params),
            rng_state=rng.internal_state,
        )

    checkpoint_dir = os.path.join(workdir, "checkpoints")

    # Setup SDEs
    sde = sde_lib.get_sde(config)
    # Add one additional round to get the exact number of samples as required.

    # num_sampling_rounds and num_bpd_rounds must be computed in all cases.
    num_sampling_rounds = int(
        math.ceil(config.eval.num_samples / config.eval.batch_size)
    )
    # Create data loaders for likelihood evaluation. Only evaluate on uniformly dequantized data
    train_ds_bpd, eval_ds_bpd = datasets.get_dataset(
        config,
        additional_dim=None,
        uniform_dequantization=True,
        evaluation=True,
        drop_last=False,
    )
    if config.eval.bpd_dataset.lower() == "train":
        ds_bpd = train_ds_bpd
    elif config.eval.bpd_dataset.lower() == "test":
        # Go over the dataset 5 times when computing likelihood on the test dataset
        ds_bpd = eval_ds_bpd
    else:
        raise ValueError(f"No bpd dataset {config.eval.bpd_dataset} recognized.")

    num_bpd_rounds = len(ds_bpd)

    if config.eval.enable_loss:
        # Build datasets
        train_ds, eval_ds = datasets.get_dataset(
            config,
            additional_dim=1,
            uniform_dequantization=config.data.uniform_dequantization,
            evaluation=True,
            drop_last=False,
        )
        # Create the one-step evaluation function when loss computation is enabled
        train_loss_fn, eval_loss_fn, state = losses.get_loss_fn(
            config, sde, score_model, state, next(rng)
        )

        ema_scales_fn = losses.get_ema_scales_fn(config)

        eval_step = losses.get_step_fn(
            eval_loss_fn,
            train=False,
            optimize_fn=optimize_fn,
            ema_scales_fn=ema_scales_fn,
        )
        # Pmap (and jit-compile) multiple evaluation steps together for faster execution
        p_eval_step = jax.pmap(
            functools.partial(jax.lax.scan, eval_step),
            axis_name="batch",
        )

    if config.eval.enable_bpd:
        # Build the likelihood computation function when likelihood is enabled
        likelihood_fn = likelihood.get_likelihood_fn(
            sde,
            score_model,
            num_repeats=5 if config.eval.bpd_dataset.lower() == "test" else 1,
        )

    # Build the sampling function when sampling is enabled
    if config.eval.enable_sampling:
        sampling_shape = (
            config.eval.batch_size // jax.local_device_count(),
            config.data.image_size,
            config.data.image_size,
            config.data.num_channels,
        )
        sampling_fn = sampling.get_sampling_fn(config, sde, score_model, sampling_shape)

    # Create different random states for different hosts in a multi-host environment (e.g., TPU pods)
    rng = hk.PRNGSequence(jax.random.fold_in(next(rng), jax.process_index()))

    # A data class for storing intermediate results to resume evaluation after pre-emption
    @flax.struct.dataclass
    class EvalMeta:
        ckpt_id: int
        sampling_round_id: int
        bpd_round_id: int
        rng_state: Any

    # Restore evaluation after pre-emption
    eval_meta = EvalMeta(
        ckpt_id=config.eval.begin_ckpt,
        sampling_round_id=-1,
        bpd_round_id=-1,
        rng_state=rng.internal_state,
    )
    eval_meta = checkpoints.restore_checkpoint(
        eval_dir, eval_meta, step=None, prefix=f"meta_{jax.process_index()}_"
    )

    # avoid not starting from config.eval.begin_ckpt.
    if eval_meta.ckpt_id < config.eval.begin_ckpt:
        eval_meta = eval_meta.replace(
            ckpt_id=config.eval.begin_ckpt,
            sampling_round_id=-1,
            bpd_round_id=-1,
            rng_state=rng.internal_state,
        )

    # Evaluation order: first loss, then likelihood, then sampling
    if eval_meta.bpd_round_id < num_bpd_rounds - 1:
        begin_ckpt = eval_meta.ckpt_id
        begin_bpd_round = eval_meta.bpd_round_id + 1
        begin_sampling_round = 0

    elif eval_meta.sampling_round_id < num_sampling_rounds - 1:
        begin_ckpt = eval_meta.ckpt_id
        begin_bpd_round = num_bpd_rounds
        begin_sampling_round = eval_meta.sampling_round_id + 1

    else:
        begin_ckpt = eval_meta.ckpt_id + 1
        begin_bpd_round = 0
        begin_sampling_round = 0

    rng.replace_internal_state(eval_meta.rng_state)
    logging.info("begin checkpoint: %d" % (begin_ckpt,))
    for ckpt in range(begin_ckpt, config.eval.end_ckpt + 1):
        ## Part 1: Load checkpoint
        # Wait if the target checkpoint doesn't exist yet
        waiting_message_printed = False
        ckpt_filename = os.path.join(checkpoint_dir, "checkpoint_{}".format(ckpt))
        while not blobfile.exists(ckpt_filename):
            if not waiting_message_printed and jax.process_index() == 0:
                logging.warning("Waiting for the arrival of checkpoint_%d" % (ckpt,))
                waiting_message_printed = True
            time.sleep(60)

        # Wait for 2 additional mins in case the file exists but is not ready for reading
        try:
            state = checkpoints.restore_checkpoint(checkpoint_dir, state, step=ckpt)
        except:
            time.sleep(60)
            try:
                state = checkpoints.restore_checkpoint(checkpoint_dir, state, step=ckpt)
            except:
                raise OSError("checkpoint file is not ready for reading")

        # Replicate the training state to prepare for pmap
        pstate = flax.jax_utils.replicate(state)
        ## Part 2: Compute loss
        if config.eval.enable_loss:
            all_losses = []
            all_log_stats = defaultdict(list)
            eval_iter = iter(eval_ds)  # pytype: disable=wrong-arg-types
            for i, batch in enumerate(eval_iter):
                eval_batch = jax.tree_util.tree_map(
                    lambda x: x.detach().cpu().numpy(), batch
                )
                next_rng = jnp.asarray(rng.take(jax.local_device_count()))
                (_, _), (
                    p_eval_loss,
                    p_eval_log_stats,
                ) = p_eval_step((next_rng, pstate), eval_batch)
                eval_loss = flax.jax_utils.unreplicate(p_eval_loss)
                eval_log_stats = flax.jax_utils.unreplicate(p_eval_log_stats)

                all_losses.extend(eval_loss)
                for key, value in eval_log_stats.items():
                    all_log_stats[key].extend(value)

                if (i + 1) % 1000 == 0 and jax.process_index() == 0:
                    logging.info("Finished %dth step loss evaluation" % (i + 1))

            # Save loss values to disk or Google Cloud Storage
            all_losses = jnp.asarray(all_losses)
            all_log_stats = jax.tree_map(lambda x: jnp.asarray(x), all_log_stats)
            with blobfile.BlobFile(
                os.path.join(eval_dir, f"ckpt_{ckpt}_loss.npz"), "wb"
            ) as fout:
                io_buffer = io.BytesIO()
                np.savez_compressed(
                    io_buffer,
                    all_losses=all_losses,
                    mean_loss=all_losses.mean(),
                    **all_log_stats,
                )
                fout.write(io_buffer.getvalue())

        ## Part 3: Compute likelihood (bits/dim)
        if config.eval.enable_bpd:
            bpds = []
            bpd_iter = iter(ds_bpd)
            for _ in range(begin_bpd_round):
                next(bpd_iter)
            for i, eval_batch in enumerate(bpd_iter):
                eval_batch = jax.tree_util.tree_map(
                    lambda x: x.detach().cpu().numpy(), eval_batch
                )
                step_rng = jnp.asarray(rng.take(jax.local_device_count()))
                bpd = likelihood_fn(step_rng, pstate, eval_batch["image"])[0]
                bpd = bpd.reshape(-1)
                bpds.extend(bpd)
                bpd_round_id = begin_bpd_round + i
                logging.info(
                    "ckpt: %d, round: %d, mean bpd: %6f"
                    % (ckpt, bpd_round_id, jnp.mean(jnp.asarray(bpds)))
                )
                # Save bits/dim to disk or Google Cloud Storage
                with blobfile.BlobFile(
                    os.path.join(
                        eval_dir,
                        f"{config.eval.bpd_dataset}_ckpt_{ckpt}_bpd_{bpd_round_id}.npz",
                    ),
                    "wb",
                ) as fout:
                    io_buffer = io.BytesIO()
                    np.savez_compressed(io_buffer, bpd)
                    fout.write(io_buffer.getvalue())

                eval_meta = eval_meta.replace(
                    ckpt_id=ckpt,
                    bpd_round_id=bpd_round_id,
                    rng_state=rng.internal_state,
                )
                # Save intermediate states to resume evaluation after pre-emption
                checkpoints.save_checkpoint(
                    eval_dir,
                    eval_meta,
                    step=ckpt * (num_bpd_rounds + num_sampling_rounds) + bpd_round_id,
                    keep=1,
                    prefix=f"meta_{jax.process_index()}_",
                )
        else:
            # Skip likelihood computation and save intermediate states for pre-emption
            eval_meta = eval_meta.replace(ckpt_id=ckpt, bpd_round_id=num_bpd_rounds - 1)
            checkpoints.save_checkpoint(
                eval_dir,
                eval_meta,
                step=ckpt * (num_bpd_rounds + num_sampling_rounds) + num_bpd_rounds - 1,
                keep=1,
                prefix=f"meta_{jax.process_index()}_",
            )

        # Generate samples and compute IS/FID/KID when enabled
        if config.eval.enable_sampling:
            logging.info(f"Start sampling evaluation for ckpt {ckpt}")
            # Run sample generation for multiple rounds to create enough samples
            # Designed to be pre-emption safe. Automatically resumes when interrupted
            for r in range(begin_sampling_round, num_sampling_rounds):
                if jax.process_index() == 0:
                    logging.info("sampling -- ckpt: %d, round: %d" % (ckpt, r))

                # Directory to save samples. Different for each host to avoid writing conflicts
                this_sample_dir = os.path.join(
                    eval_dir, f"ckpt_{ckpt}_host_{jax.process_index()}"
                )
                blobfile.makedirs(this_sample_dir)
                sample_rng = jnp.asarray(rng.take(jax.local_device_count()))
                samples, n = sampling_fn(sample_rng, pstate)
                samples = (samples + 1.0) / 2.0
                samples = np.clip(samples * 255.0, 0, 255).astype(np.uint8)
                samples = samples.reshape(
                    (
                        -1,
                        config.data.image_size,
                        config.data.image_size,
                        config.data.num_channels,
                    )
                )
                # Write samples to disk or Google Cloud Storage
                with blobfile.BlobFile(
                    os.path.join(this_sample_dir, f"samples_{r}.npz"),
                    "wb",
                ) as fout:
                    io_buffer = io.BytesIO()
                    np.savez_compressed(io_buffer, samples=samples)
                    fout.write(io_buffer.getvalue())

                # Save image samples and submit to the FID evaluation website
                if r == num_sampling_rounds - 1:
                    # Collect samples from all hosts and sampling rounds
                    if jax.process_index() == 0:
                        all_samples = get_samples_from_ckpt(eval_dir, ckpt)
                    all_samples = all_samples[: config.eval.num_samples]
                    sample_path = os.path.join(eval_dir, f"ckpt_{ckpt}_samples.npz")
                    with blobfile.BlobFile(sample_path, "wb") as fout:
                        io_buffer = io.BytesIO()
                        np.savez_compressed(io_buffer, all_samples)
                        fout.write(io_buffer.getvalue())

                # Update the intermediate evaluation state
                eval_meta = eval_meta.replace(
                    ckpt_id=ckpt, sampling_round_id=r, rng_state=rng.internal_state
                )
                # Save intermediate states to resume evaluation after pre-emption
                checkpoints.save_checkpoint(
                    eval_dir,
                    eval_meta,
                    step=ckpt * (num_sampling_rounds + num_bpd_rounds)
                    + r
                    + num_bpd_rounds,
                    keep=1,
                    prefix=f"meta_{jax.process_index()}_",
                )

        else:
            # Skip sampling and save intermediate evaluation states for pre-emption
            eval_meta = eval_meta.replace(
                ckpt_id=ckpt,
                sampling_round_id=num_sampling_rounds - 1,
                rng_state=rng.internal_state,
            )
            checkpoints.save_checkpoint(
                eval_dir,
                eval_meta,
                step=ckpt * (num_sampling_rounds + num_bpd_rounds)
                + num_sampling_rounds
                - 1
                + num_bpd_rounds,
                keep=1,
                prefix=f"meta_{jax.process_index()}_",
            )

        begin_bpd_round = 0
        begin_sampling_round = 0

    # Remove all meta files after finishing evaluation
    meta_files = blobfile.glob(os.path.join(eval_dir, f"meta_{jax.process_index()}_*"))
    for file in meta_files:
        blobfile.remove(file)