in jcm/evaluate.py [0:0]
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)