run_lib_label.py

# coding=utf-8
# Copyright 2020 The Google Research Authors.
#
# 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.

# pylint: skip-file
"""Training and evaluation for score-based generative models. """

import gc
import io
import os
import time

import numpy as np
import tensorflow as tf
# import tensorflow_gan as tfgan
import logging
# Keep the import below for registering all model definitions
from models import ddpm
import losses
import sampling, sampling_multivariate, sampling_conditional, sampling_latent, sampling_label
from models import utils as mutils
from models.ema import ExponentialMovingAverage
import datasets
# import evaluation
import likelihood
import sde_lib
from absl import flags
import torch
from torch.utils import tensorboard
from torchvision.utils import make_grid, save_image
from utils import save_checkpoint, restore_checkpoint
from scipy.fft import dct

FLAGS = flags.FLAGS

def compute_dct(n=10):
  #computes a matrix of size n*n such that it can be used to perform discrete cosine transform
  # note that axis =1 means that the columns are the basis, i.e fist col has all the same values

  D = dct(np.eye(n), axis =1)
  W = torch.from_numpy(D).float()
  return W

def tensor_from_idx(eval_label, label_size = 10):
    m = eval_label.shape[0]
    label = torch.zeros(m,label_size)
    label[torch.arange(m), eval_label] = 1
    return label


def plot_samples(sample_dir, sample, step):
  this_sample_dir = os.path.join(sample_dir, "iter_{}".format(step))
  tf.io.gfile.makedirs(this_sample_dir)
  nrow = int(np.sqrt(sample.shape[0]))
  image_grid = make_grid(sample, nrow, padding=2)
  sample = np.clip(sample.permute(0, 2, 3, 1).cpu().numpy() * 255, 0, 255).astype(np.uint8)
  with tf.io.gfile.GFile(
      os.path.join(this_sample_dir, "sample.np"), "wb") as fout:
    np.save(fout, sample)

  with tf.io.gfile.GFile(
      os.path.join(this_sample_dir, "sample.png"), "wb") as fout:
    save_image(image_grid, fout)

def plot_samples_analysis(sample_dir, sample_stack, step, pos_index):
  this_sample_dir = os.path.join(sample_dir, "iter_{}".format(step), str(pos_index))
  tf.io.gfile.makedirs(this_sample_dir)
  for k in range(sample_stack.shape[-1]):
    sample = sample_stack[:,:,:,:,k]
    nrow = int(np.sqrt(sample.shape[0]))
    image_grid = make_grid(sample, nrow, padding=2)
    sample = np.clip(sample.permute(0, 2, 3, 1).cpu().numpy() * 255, 0, 255).astype(np.uint8)
    # with tf.io.gfile.GFile(
    #     os.path.join(this_sample_dir, "sample.np"), "wb") as fout:
    #   np.save(fout, sample)
   
    image_name = "sample_"+str(k)+".png"
    
    with tf.io.gfile.GFile(
        os.path.join(this_sample_dir, image_name), "wb") as fout:
      save_image(image_grid, fout)

def plot_samples_conditional(sample_dir, sample_stack, step, class_n):
  this_sample_dir = os.path.join(sample_dir, "iter_{}".format(step))
  tf.io.gfile.makedirs(this_sample_dir)
  for k in range(sample_stack.shape[-1]):
    sample = sample_stack[:,:,:,:,k]
    nrow = int(np.sqrt(sample.shape[0]))
    image_grid = make_grid(sample, nrow, padding=2)
    sample = np.clip(sample.permute(0, 2, 3, 1).cpu().numpy() * 255, 0, 255).astype(np.uint8)
    # with tf.io.gfile.GFile(
    #     os.path.join(this_sample_dir, "sample.np"), "wb") as fout:
    #   np.save(fout, sample)
    if len(class_n)>1:
      image_name = "sample_"+str(class_n[k])+".png"
    else:
      
      image_name = "sample_"+str(class_n[0])+".png"

    with tf.io.gfile.GFile(
        os.path.join(this_sample_dir, image_name), "wb") as fout:
      save_image(image_grid, fout)

def plot_scores(sample_dir, sample, step):
  this_sample_dir = os.path.join(sample_dir, "iter_{}".format(step))
  tf.io.gfile.makedirs(this_sample_dir)
  nrow = sample.shape[-1]
  sample = sample.permute(0,4,1,2,3)
  sample = torch.abs(sample.contiguous().view(sample.shape[0]*sample.shape[1], sample.shape[2], sample.shape[3], sample.shape[4]))
  sample = sample/torch.max(sample)
  image_grid = make_grid(sample, nrow, padding=2)
  sample = np.clip(sample.permute(0, 2, 3, 1).cpu().numpy() * 255, 0, 255).astype(np.uint8)
  
  with tf.io.gfile.GFile(
      os.path.join(this_sample_dir, "sample.np"), "wb") as fout:
    np.save(fout, sample)

  with tf.io.gfile.GFile(
      os.path.join(this_sample_dir, "sample.png"), "wb") as fout:
    save_image(image_grid, fout)
 
def plot_samples_multivariate(sample_dir, sample, sample_stack, step):
  k = sample_stack.shape[-1]
  this_sample_dir = os.path.join(sample_dir, "iter_{}".format(step))
  tf.io.gfile.makedirs(this_sample_dir)

  nrow = int(np.sqrt(sample.shape[0]))
  image_grid = make_grid(sample, nrow, padding=2)
  sample = np.clip(sample.permute(0, 2, 3, 1).cpu().numpy() * 255, 0, 255).astype(np.uint8)
  with tf.io.gfile.GFile(
      os.path.join(this_sample_dir, "sample.np"), "wb") as fout:
    np.save(fout, sample)

  with tf.io.gfile.GFile(
      os.path.join(this_sample_dir, "sample.png"), "wb") as fout:
    save_image(image_grid, fout)

  for i in range(k):
    sample_i = sample_stack[:,:,:,:,i]
    nrow = int(np.sqrt(sample_i.shape[0]))
    image_grid = make_grid(sample_i, nrow, padding=2)
    sample = np.clip(sample_i.permute(0, 2, 3, 1).cpu().numpy() * 255, 0, 255).astype(np.uint8)
    img_name = "sample_"+str(i)
    with tf.io.gfile.GFile(
        os.path.join(this_sample_dir, img_name+".np"), "wb") as fout:
      np.save(fout, sample)

    with tf.io.gfile.GFile(
        os.path.join(this_sample_dir, img_name+".png"), "wb") as fout:
      save_image(image_grid, fout)


def train(config, workdir):
  """Runs the training pipeline.

  Args:
    config: Configuration to use.
    workdir: Working directory for checkpoints and TF summaries. If this
      contains checkpoint training will be resumed from the latest checkpoint.
  """

  # Create directories for experimental logs
  sample_dir = os.path.join(workdir, "samples")
  tf.io.gfile.makedirs(sample_dir)

  tb_dir = os.path.join(workdir, "tensorboard")
  tf.io.gfile.makedirs(tb_dir)
  writer = tensorboard.SummaryWriter(tb_dir)

  # Initialize model.
  score_model = mutils.create_model(config)
  ema = ExponentialMovingAverage(score_model.parameters(), decay=config.model.ema_rate)
  optimizer = losses.get_optimizer(config, score_model.parameters())
  state = dict(optimizer=optimizer, model=score_model, ema=ema, step=0)

  # Create checkpoints directory
  checkpoint_dir = os.path.join(workdir, "checkpoints")
  # Intermediate checkpoints to resume training after pre-emption in cloud environments
  checkpoint_meta_dir = os.path.join(workdir, "checkpoints-meta", "checkpoint.pth")
  tf.io.gfile.makedirs(checkpoint_dir)
  tf.io.gfile.makedirs(os.path.dirname(checkpoint_meta_dir))
  # Resume training when intermediate checkpoints are detected
  state = restore_checkpoint(checkpoint_meta_dir, state, config.device)
  initial_step = int(state['step'])

  # Build data iterators
  train_ds, eval_ds, _ = datasets.get_dataset(config,
                                              uniform_dequantization=config.data.uniform_dequantization)
  train_iter = iter(train_ds)  # pytype: disable=wrong-arg-types
  eval_iter = iter(eval_ds)  # pytype: disable=wrong-arg-types
  # Create data normalizer and its inverse
  scaler = datasets.get_data_scaler(config)
  inverse_scaler = datasets.get_data_inverse_scaler(config)

  # Setup SDEs
  if config.training.sde.lower() == 'vpsde':
    sde = sde_lib.VPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
    sampling_eps = 1e-3
  elif config.training.sde.lower() == 'subvpsde':
    sde = sde_lib.subVPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
    sampling_eps = 1e-3
  elif config.training.sde.lower() == 'vesde':
    sde = sde_lib.VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
    sampling_eps = 1e-5
  else:
    raise NotImplementedError(f"SDE {config.training.sde} unknown.")

  # Build one-step training and evaluation functions
  optimize_fn = losses.optimization_manager(config)
  continuous = config.training.continuous
  reduce_mean = config.training.reduce_mean
  likelihood_weighting = config.training.likelihood_weighting
  train_step_fn = losses.get_step_fn(sde, train=True, optimize_fn=optimize_fn,
                                     reduce_mean=reduce_mean, continuous=continuous,
                                     likelihood_weighting=likelihood_weighting, config = config)
  eval_step_fn = losses.get_step_fn(sde, train=False, optimize_fn=optimize_fn,
                                    reduce_mean=reduce_mean, continuous=continuous,
                                    likelihood_weighting=likelihood_weighting, config = config)

  # Building sampling functions
  if config.training.snapshot_sampling:
    sampling_shape = (config.training.batch_size, config.data.num_channels,
                      config.data.image_size, config.data.image_size)

    if config.training.conditional_model == 'label':
      sampling_fn = sampling_latent.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)
    elif config.sampling.type == 'multivariate':
      sampling_fn = sampling_multivariate.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)
    elif config.sampling.type == 'conditional':
      sampling_fn = sampling_conditional.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)
    else:
      sampling_fn = sampling.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)

  num_train_steps = config.training.n_iters

  W = compute_dct(n=10).to(config.device)

  # In case there are multiple hosts (e.g., TPU pods), only log to host 0
  logging.info("Starting training loop at step %d." % (initial_step,))

  for step in range(initial_step, num_train_steps + 1):
    # Convert data to JAX arrays and normalize them. Use ._numpy() to avoid copy.
    data = next(train_iter)
    batch = torch.from_numpy(data['image']._numpy()).to(config.device).float()
    batch = batch.permute(0, 3, 1, 2)
    batch = scaler(batch)

    label = torch.from_numpy(data['label']._numpy())
    label = tensor_from_idx(label).to(config.device)
    label = torch.mm(label,W)
    # Execute one training step
    loss = train_step_fn(state, batch, label)
    if step % config.training.log_freq == 0:
      logging.info("step: %d, training_loss: %.5e" % (step, loss.item()))
      writer.add_scalar("training_loss", loss, step)

    # Save a temporary checkpoint to resume training after pre-emption periodically
    if step != 0 and step % config.training.snapshot_freq_for_preemption == 0:
      save_checkpoint(checkpoint_meta_dir, state)

    # Report the loss on an evaluation dataset periodically
    if step % config.training.eval_freq == 0:
      eval_data = next(eval_iter)
      eval_batch = torch.from_numpy(eval_data['image']._numpy()).to(config.device).float()
      eval_batch = eval_batch.permute(0, 3, 1, 2)
      eval_batch = scaler(eval_batch)

      eval_label = torch.from_numpy(eval_data['label']._numpy())
      eval_label = tensor_from_idx(eval_label).to(config.device)
      eval_label = torch.mm(eval_label, W)

      eval_loss = eval_step_fn(state, eval_batch, eval_label)
      logging.info("step: %d, eval_loss: %.5e" % (step, eval_loss.item()))
      writer.add_scalar("eval_loss", eval_loss.item(), step)

    # Save a checkpoint periodically and generate samples if needed
    if step != 0 and step % config.training.snapshot_freq == 0 or step == num_train_steps:
      # Save the checkpoint.
      save_step = step // config.training.snapshot_freq
      save_checkpoint(os.path.join(checkpoint_dir, f'checkpoint_{save_step}.pth'), state)

      # Generate and save samples
      if config.training.snapshot_sampling:
        ema.store(score_model.parameters())
        ema.copy_to(score_model.parameters())
        if config.training.conditional_model == 'label':
          class_n = save_step%10
          latent = torch.zeros(10)
          latent[class_n] = 1

          # class_n_2 = (save_step+1)%10
          # latent_2 = torch.zeros(10)
          # latent_2[class_n_2] = 1

          latent_stack = latent.unsqueeze(-1).to(config.device)
          latent_stack = torch.mm(W.T, latent_stack).squeeze()
          sample, n = sampling_fn(score_model, latent_stack)
        elif config.sampling.type == 'multivariate':
          sample, sample_stack, n = sampling_fn(score_model)
        elif config.sampling.type == 'conditional':
          class_n = save_step%10
          latent = torch.zeros(10)
          latent[class_n] = 1
          sample, n = sampling_fn(score_model, latent)
        else:
          sample, n = sampling_fn(score_model)
        ema.restore(score_model.parameters())
        if config.sampling.type == 'multivariate':
          plot_samples_multivariate(sample_dir, sample, sample_stack, step) #either single or multivariate
        elif config.sampling.type == 'conditional' or config.training.conditional_model == 'label':
          plot_samples_conditional(sample_dir, sample.unsqueeze(-1), step, [class_n])
        else:
          plot_samples(sample_dir, sample, step) #either single or multivariate

        
def visualize(config, workdir, visualization_folder="viz"):
  # Create directory to eval_folder
  viz_dir = os.path.join(workdir, visualization_folder)
  tf.io.gfile.makedirs(viz_dir)

  # Build data pipeline
  train_ds, eval_ds, _ = datasets.get_dataset(config,
                                              uniform_dequantization=config.data.uniform_dequantization,
                                              evaluation=True)

  # Create data normalizer and its inverse
  scaler = datasets.get_data_scaler(config)
  inverse_scaler = datasets.get_data_inverse_scaler(config)

  # Initialize model
  score_model = mutils.create_model(config)
  optimizer = losses.get_optimizer(config, score_model.parameters())
  ema = ExponentialMovingAverage(score_model.parameters(), decay=config.model.ema_rate)
  state = dict(optimizer=optimizer, model=score_model, ema=ema, step=0)

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

  # Setup SDEs
  if config.training.sde.lower() == 'vpsde':
    sde = sde_lib.VPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
    sampling_eps = 1e-3
  elif config.training.sde.lower() == 'subvpsde':
    sde = sde_lib.subVPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
    sampling_eps = 1e-3
  elif config.training.sde.lower() == 'vesde':
    sde = sde_lib.VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
    sampling_eps = 1e-5
  else:
    raise NotImplementedError(f"SDE {config.training.sde} unknown.")

  # Build the sampling function when sampling is enabled

  sampling_shape = (config.eval.batch_size,
                    config.data.num_channels,
                    config.data.image_size, config.data.image_size)
  if config.training.conditional_model == 'equal_energy':
      sampling_fn = sampling_equal_energy.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)
  elif config.sampling.type == 'multivariate':
    sampling_fn = sampling_multivariate.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)
  else:
    sampling_fn = sampling.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)


  begin_ckpt = config.eval.begin_ckpt
  logging.info("begin checkpoint: %d" % (begin_ckpt,))
  ckpt = begin_ckpt
  # for ckpt in range(begin_ckpt, config.eval.end_ckpt + 1):
  # Wait if the target checkpoint doesn't exist yet
  # waiting_message_printed = False
  # ckpt_filename = os.path.join(checkpoint_dir, "checkpoint_{}.pth".format(ckpt))
  # while not tf.io.gfile.exists(ckpt_filename):
  #   if not waiting_message_printed:
  #     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
  ckpt_path = os.path.join(checkpoint_dir, f'checkpoint_{ckpt}.pth')
  try:
    state = restore_checkpoint(ckpt_path, state, device=config.device)
  except:
    time.sleep(60)
    try:
      state = restore_checkpoint(ckpt_path, state, device=config.device)
    except:
      time.sleep(120)
      state = restore_checkpoint(ckpt_path, state, device=config.device)
  ema.copy_to(score_model.parameters())

  sample_dir = os.path.join(viz_dir, "sample")
  score_dir = os.path.join(viz_dir, "score_plot")

  if config.training.conditional_model == 'equal_energy':
    class_n = 0
    latent = torch.zeros(3)
    latent[class_n] = 0.5

    class_n_2 = 0
    latent_2 = torch.zeros(3)
    latent_2[class_n_2] = 0.7

    latent_stack = torch.cat((latent.unsqueeze(-1), latent_2.unsqueeze(-1)), dim=-1)
    sample, n = sampling_fn(score_model, latent_stack)
    plot_samples_conditional(sample_dir, sample, 12.0, [0.5,0.7])

  elif config.sampling.type == 'multivariate' and config.sampling.plot_score:
    sample, sample_stack, n, score_stack = sampling_fn(score_model, return_score='score')
    plot_scores(score_dir, score_stack, ckpt)
    plot_samples_multivariate(sample_dir, sample, sample_stack, ckpt)

  elif config.sampling.type == 'multivariate':
    sample, sample_stack, n = sampling_fn(score_model)
    plot_samples_multivariate(sample_dir, sample, sample_stack, ckpt)
  elif config.sampling.type == 'conditional':
    sample, n = sampling_fn(score_model)
    plot_samples_conditional(sample_dir, sample, ckpt)
  else:
    sample, n = sampling_fn(score_model)
    plot_samples(sample_dir, sample, ckpt)
  ema.restore(score_model.parameters())


def analyze(config, workdir, visualization_folder="viz"):
  # Create directory to eval_folder
  viz_dir = os.path.join(workdir, visualization_folder)
  tf.io.gfile.makedirs(viz_dir)

  # Build data pipeline
  train_ds, eval_ds, _ = datasets.get_dataset(config,
                                              uniform_dequantization=config.data.uniform_dequantization,
                                              evaluation=True)

  # Create data normalizer and its inverse
  scaler = datasets.get_data_scaler(config)
  inverse_scaler = datasets.get_data_inverse_scaler(config)

  # Initialize model
  score_model = mutils.create_model(config)
  optimizer = losses.get_optimizer(config, score_model.parameters())
  ema = ExponentialMovingAverage(score_model.parameters(), decay=config.model.ema_rate)
  state = dict(optimizer=optimizer, model=score_model, ema=ema, step=0)

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

  # Setup SDEs
  if config.training.sde.lower() == 'vpsde':
    sde = sde_lib.VPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
    sampling_eps = 1e-3
  elif config.training.sde.lower() == 'subvpsde':
    sde = sde_lib.subVPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
    sampling_eps = 1e-3
  elif config.training.sde.lower() == 'vesde':
    sde = sde_lib.VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
    sampling_eps = 1e-5
  else:
    raise NotImplementedError(f"SDE {config.training.sde} unknown.")

  # Build the sampling function when sampling is enabled

  sampling_shape = (config.eval.batch_size,
                    config.data.num_channels,
                    config.data.image_size, config.data.image_size)
  if config.training.conditional_model == 'label':
    analysis_fn = sampling_label.get_analysis_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)
  # if config.training.conditional_model == 'latent_factor':
  #   analysis_fn = sampling_latent_factor.get_analysis_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps, mode = 'interpolate')
  
  # else:
  #   sampling_fn = sampling.get_sampling_fn(config, sde, sampling_shape, inverse_scaler, sampling_eps)
  
  train_ds, eval_ds, _ = datasets.get_dataset(config,
                                              uniform_dequantization=config.data.uniform_dequantization,
                                              evaluation=True)


  begin_ckpt = config.eval.begin_ckpt
  logging.info("begin checkpoint: %d" % (begin_ckpt,))
  ckpt = begin_ckpt
  # for ckpt in range(begin_ckpt, config.eval.end_ckpt + 1):
  # Wait if the target checkpoint doesn't exist yet
  # waiting_message_printed = False
  # ckpt_filename = os.path.join(checkpoint_dir, "checkpoint_{}.pth".format(ckpt))
  # while not tf.io.gfile.exists(ckpt_filename):
  #   if not waiting_message_printed:
  #     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
  ckpt_path = os.path.join(checkpoint_dir, f'checkpoint_{ckpt}.pth')
  try:
    state = restore_checkpoint(ckpt_path, state, device=config.device)
  except:
    time.sleep(60)
    try:
      state = restore_checkpoint(ckpt_path, state, device=config.device)
    except:
      time.sleep(120)
      state = restore_checkpoint(ckpt_path, state, device=config.device)
  ema.copy_to(score_model.parameters())

  sample_dir = os.path.join(viz_dir, "analyze")
  # score_dir = os.path.join(viz_dir, "score_plot")
  pos_index = 2
  
  # step = (torch.FloatTensor([ckpt]) + 0.1*torch.FloatTensor([class_n])).item()
  step = ckpt
  eval_iter = iter(eval_ds)  # pytype: disable=wrong-arg-types
  train_iter = iter(train_ds)
  # for i, batch in enumerate(train_iter):
  batch = next(train_iter)
  eval_batch = torch.from_numpy(batch['image']._numpy()).to(config.device).float()
  eval_batch = eval_batch.permute(0, 3, 1, 2)
  eval_batch_model = scaler(eval_batch)

  W = compute_dct(n=10).to(config.device)

  i,j= 4,7
  z1 = torch.zeros(10)
  z1[i] = 1

  z2 = torch.zeros(10)
  z2[j] = 1

  # class_n_2 = (save_step+1)%10
  # latent_2 = torch.zeros(10)
  # latent_2[class_n_2] = 1

  latent_stack = torch.cat((z1.unsqueeze(-1), z2.unsqueeze(-1)), dim=-1).to(config.device)
  latent_stack = torch.mm(W.T, latent_stack).squeeze()

  # if config.training.conditional_model == 'latent_factor':
  #   pos_index = [5,6]
  #   samples = analysis_fn(score_model, eval_batch_model, pos_index)
  #   # plot_samples_analysis(sample_dir, samples, step, pos_index)

  # else:
  sample1, sample2, sample3 = analysis_fn(score_model, latent_stack)
  samples = torch.cat((eval_batch.unsqueeze(-1), sample1.unsqueeze(-1), sample2.unsqueeze(-1), sample3.unsqueeze(-1)), dim =-1)
  # samples=torch.cat((eval_batch.unsqueeze(-1),samples), dim=-1)
  plot_samples_analysis(sample_dir, samples, step, str(i)+','+str(j))
        # plot_samples_conditional(sample_dir, sample, 12.0, [0.5,0.7])
  print("Done!")

# 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)
#   tf.io.gfile.makedirs(eval_dir)

#   # Build data pipeline
#   train_ds, eval_ds, _ = datasets.get_dataset(config,
#                                               uniform_dequantization=config.data.uniform_dequantization,
#                                               evaluation=True)

#   # Create data normalizer and its inverse
#   scaler = datasets.get_data_scaler(config)
#   inverse_scaler = datasets.get_data_inverse_scaler(config)

#   # Initialize model
#   score_model = mutils.create_model(config)
#   optimizer = losses.get_optimizer(config, score_model.parameters())
#   ema = ExponentialMovingAverage(score_model.parameters(), decay=config.model.ema_rate)
#   state = dict(optimizer=optimizer, model=score_model, ema=ema, step=0)

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

#   # Setup SDEs
#   if config.training.sde.lower() == 'vpsde':
#     sde = sde_lib.VPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
#     sampling_eps = 1e-3
#   elif config.training.sde.lower() == 'subvpsde':
#     sde = sde_lib.subVPSDE(beta_min=config.model.beta_min, beta_max=config.model.beta_max, N=config.model.num_scales)
#     sampling_eps = 1e-3
#   elif config.training.sde.lower() == 'vesde':
#     sde = sde_lib.VESDE(sigma_min=config.model.sigma_min, sigma_max=config.model.sigma_max, N=config.model.num_scales)
#     sampling_eps = 1e-5
#   else:
#     raise NotImplementedError(f"SDE {config.training.sde} unknown.")

#   # Create the one-step evaluation function when loss computation is enabled
#   if config.eval.enable_loss:
#     optimize_fn = losses.optimization_manager(config)
#     continuous = config.training.continuous
#     likelihood_weighting = config.training.likelihood_weighting

#     reduce_mean = config.training.reduce_mean
#     eval_step = losses.get_step_fn(sde, train=False, optimize_fn=optimize_fn,
#                                    reduce_mean=reduce_mean,
#                                    continuous=continuous,
#                                    likelihood_weighting=likelihood_weighting)


#   # Create data loaders for likelihood evaluation. Only evaluate on uniformly dequantized data
#   train_ds_bpd, eval_ds_bpd, _ = datasets.get_dataset(config,
#                                                       uniform_dequantization=True, evaluation=True)
#   if config.eval.bpd_dataset.lower() == 'train':
#     ds_bpd = train_ds_bpd
#     bpd_num_repeats = 1
#   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
#     bpd_num_repeats = 5
#   else:
#     raise ValueError(f"No bpd dataset {config.eval.bpd_dataset} recognized.")

#   # Build the likelihood computation function when likelihood is enabled
#   if config.eval.enable_bpd:
#     likelihood_fn = likelihood.get_likelihood_fn(sde, inverse_scaler)

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

#   # Use inceptionV3 for images with resolution higher than 256.
#   inceptionv3 = config.data.image_size >= 256
#   inception_model = evaluation.get_inception_model(inceptionv3=inceptionv3)

#   begin_ckpt = config.eval.begin_ckpt
#   logging.info("begin checkpoint: %d" % (begin_ckpt,))
#   for ckpt in range(begin_ckpt, config.eval.end_ckpt + 1):
#     # Wait if the target checkpoint doesn't exist yet
#     waiting_message_printed = False
#     ckpt_filename = os.path.join(checkpoint_dir, "checkpoint_{}.pth".format(ckpt))
#     while not tf.io.gfile.exists(ckpt_filename):
#       if not waiting_message_printed:
#         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
#     ckpt_path = os.path.join(checkpoint_dir, f'checkpoint_{ckpt}.pth')
#     try:
#       state = restore_checkpoint(ckpt_path, state, device=config.device)
#     except:
#       time.sleep(60)
#       try:
#         state = restore_checkpoint(ckpt_path, state, device=config.device)
#       except:
#         time.sleep(120)
#         state = restore_checkpoint(ckpt_path, state, device=config.device)
#     ema.copy_to(score_model.parameters())
#     # Compute the loss function on the full evaluation dataset if loss computation is enabled
#     if config.eval.enable_loss:
#       all_losses = []
#       eval_iter = iter(eval_ds)  # pytype: disable=wrong-arg-types
#       for i, batch in enumerate(eval_iter):
#         eval_batch = torch.from_numpy(batch['image']._numpy()).to(config.device).float()
#         eval_batch = eval_batch.permute(0, 3, 1, 2)
#         eval_batch = scaler(eval_batch)
#         eval_loss = eval_step(state, eval_batch)
#         all_losses.append(eval_loss.item())
#         if (i + 1) % 1000 == 0:
#           logging.info("Finished %dth step loss evaluation" % (i + 1))

#       # Save loss values to disk or Google Cloud Storage
#       all_losses = np.asarray(all_losses)
#       with tf.io.gfile.GFile(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())
#         fout.write(io_buffer.getvalue())

#     # Compute log-likelihoods (bits/dim) if enabled
#     if config.eval.enable_bpd:
#       bpds = []
#       for repeat in range(bpd_num_repeats):
#         bpd_iter = iter(ds_bpd)  # pytype: disable=wrong-arg-types
#         for batch_id in range(len(ds_bpd)):
#           batch = next(bpd_iter)
#           eval_batch = torch.from_numpy(batch['image']._numpy()).to(config.device).float()
#           eval_batch = eval_batch.permute(0, 3, 1, 2)
#           eval_batch = scaler(eval_batch)
#           bpd = likelihood_fn(score_model, eval_batch)[0]
#           bpd = bpd.detach().cpu().numpy().reshape(-1)
#           bpds.extend(bpd)
#           logging.info(
#             "ckpt: %d, repeat: %d, batch: %d, mean bpd: %6f" % (ckpt, repeat, batch_id, np.mean(np.asarray(bpds))))
#           bpd_round_id = batch_id + len(ds_bpd) * repeat
#           # Save bits/dim to disk or Google Cloud Storage
#           with tf.io.gfile.GFile(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())

#     # Generate samples and compute IS/FID/KID when enabled
#     if config.eval.enable_sampling:
#       num_sampling_rounds = config.eval.num_samples // config.eval.batch_size + 1
#       for r in range(num_sampling_rounds):
#         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}")
#         tf.io.gfile.makedirs(this_sample_dir)
#         samples, n = sampling_fn(score_model)
#         samples = np.clip(samples.permute(0, 2, 3, 1).cpu().numpy() * 255., 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 tf.io.gfile.GFile(
#             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())

#         # Force garbage collection before calling TensorFlow code for Inception network
#         gc.collect()
#         latents = evaluation.run_inception_distributed(samples, inception_model,
#                                                        inceptionv3=inceptionv3)
#         # Force garbage collection again before returning to JAX code
#         gc.collect()
#         # Save latent represents of the Inception network to disk or Google Cloud Storage
#         with tf.io.gfile.GFile(
#             os.path.join(this_sample_dir, f"statistics_{r}.npz"), "wb") as fout:
#           io_buffer = io.BytesIO()
#           np.savez_compressed(
#             io_buffer, pool_3=latents["pool_3"], logits=latents["logits"])
#           fout.write(io_buffer.getvalue())

#       # Compute inception scores, FIDs and KIDs.
#       # Load all statistics that have been previously computed and saved for each host
#       all_logits = []
#       all_pools = []
#       this_sample_dir = os.path.join(eval_dir, f"ckpt_{ckpt}")
#       stats = tf.io.gfile.glob(os.path.join(this_sample_dir, "statistics_*.npz"))
#       for stat_file in stats:
#         with tf.io.gfile.GFile(stat_file, "rb") as fin:
#           stat = np.load(fin)
#           if not inceptionv3:
#             all_logits.append(stat["logits"])
#           all_pools.append(stat["pool_3"])

#       if not inceptionv3:
#         all_logits = np.concatenate(all_logits, axis=0)[:config.eval.num_samples]
#       all_pools = np.concatenate(all_pools, axis=0)[:config.eval.num_samples]

#       # Load pre-computed dataset statistics.
#       data_stats = evaluation.load_dataset_stats(config)
#       data_pools = data_stats["pool_3"]

#       # Compute FID/KID/IS on all samples together.
#       if not inceptionv3:
#         inception_score = tfgan.eval.classifier_score_from_logits(all_logits)
#       else:
#         inception_score = -1

#       fid = tfgan.eval.frechet_classifier_distance_from_activations(
#         data_pools, all_pools)
#       # Hack to get tfgan KID work for eager execution.
#       tf_data_pools = tf.convert_to_tensor(data_pools)
#       tf_all_pools = tf.convert_to_tensor(all_pools)
#       kid = tfgan.eval.kernel_classifier_distance_from_activations(
#         tf_data_pools, tf_all_pools).numpy()
#       del tf_data_pools, tf_all_pools

#       logging.info(
#         "ckpt-%d --- inception_score: %.6e, FID: %.6e, KID: %.6e" % (
#           ckpt, inception_score, fid, kid))

#       with tf.io.gfile.GFile(os.path.join(eval_dir, f"report_{ckpt}.npz"),
#                              "wb") as f:
#         io_buffer = io.BytesIO()
#         np.savez_compressed(io_buffer, IS=inception_score, fid=fid, kid=kid)
#         f.write(io_buffer.getvalue())