Method __init__

diffusion_tf/tpu_utils/tpu_utils.py:256–308

(self, tpu_name, model_constructor, total_bs, dataset, inception_bs=8, num_inception_samples=1024, limit_dataset_size=0)

class EvalWorker:
  def __init__(self, tpu_name, model_constructor, total_bs, dataset, inception_bs=8, num_inception_samples=1024, limit_dataset_size=0):

    self.resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu=tpu_name)
    tf.tpu.experimental.initialize_tpu_system(self.resolver)
    self.strategy = tf.distribute.experimental.TPUStrategy(self.resolver)

    self.num_cores = self.strategy.num_replicas_in_sync
    assert total_bs % self.num_cores == 0
    self.total_bs = total_bs
    self.local_bs = total_bs // self.num_cores
    print('num cores: {}'.format(self.num_cores))
    print('total batch size: {}'.format(self.total_bs))
    print('local batch size: {}'.format(self.local_bs))
    self.num_inception_samples = num_inception_samples
    assert inception_bs % self.num_cores == 0
    self.inception_bs = inception_bs
    self.inception_local_bs = inception_bs // self.num_cores
    self.dataset = dataset
    assert dataset.num_classes == 1, 'not supported'

    # TPU context
    with self.strategy.scope():
      # Inception network on real data
      print('===== INCEPTION =====')
      # Eval dataset iterator (this is the training set without repeat & shuffling)
      self.inception_real_train = InceptionFeatures(
        dataset=dataset.train_one_pass_input_fn(params={'batch_size': total_bs}), strategy=self.strategy, limit_dataset_size=limit_dataset_size // total_bs)
      # Val dataset, if it exists
      val_ds = dataset.eval_input_fn(params={'batch_size': total_bs})
      self.inception_real_val = None if val_ds is None else InceptionFeatures(dataset=val_ds, strategy=self.strategy, limit_dataset_size=limit_dataset_size // total_bs)

      img_batch_shape = self.inception_real_train.ds_iterator.output_shapes['image'].as_list()
      assert img_batch_shape[0] == self.local_bs

      # Model
      self.model = model_constructor()
      assert isinstance(self.model, Model)

      # Eval/samples graphs
      print('===== SAMPLES =====')
      self.samples_outputs, self.samples_inception = self._make_sampling_graph(
        img_shape=img_batch_shape[1:], with_inception=True)

      # Model with EMA parameters
      self.global_step = tf.train.get_or_create_global_step()
      print('===== EMA =====')
      ema, _ = make_ema(global_step=self.global_step, ema_decay=1e-10, trainable_variables=tf.trainable_variables())

      # EMA versions of the above
      with utils.ema_scope(ema):
        print('===== EMA SAMPLES =====')
        self.ema_samples_outputs, self.ema_samples_inception = self._make_sampling_graph(
          img_shape=img_batch_shape[1:], with_inception=True)

  def _make_sampling_graph(self, img_shape, with_inception):