Function model_fn

diffusion_tf/tpu_utils/tpu_utils.py:130–180

(features, params, mode)

  # model_fn for TPUEstimator
  def model_fn(features, params, mode):
    local_bs = params['batch_size']
    print('Global batch size: {}, local batch size: {}'.format(total_bs, local_bs))
    assert total_bs == num_tpu_replicas() * local_bs

    assert mode == tf.estimator.ModeKeys.TRAIN, 'only TRAIN mode supported'
    assert features['image'].shape[0] == local_bs
    assert features['label'].shape == [local_bs] and features['label'].dtype == tf.int32
    # assert labels.dtype == features['label'].dtype and labels.shape == features['label'].shape

    del params

    ##########

    # create model
    model = model_constructor()
    assert isinstance(model, Model)

    # training loss
    train_info_dict = model.train_fn(normalize_data(tf.cast(features['image'], tf.float32)), features['label'])
    loss = train_info_dict['loss']
    assert loss.shape == []

    # train op
    trainable_variables = tf.trainable_variables()
    print('num params: {:,}'.format(sum(int(np.prod(p.shape.as_list())) for p in trainable_variables)))
    global_step = tf.train.get_or_create_global_step()
    warmed_up_lr = utils.get_warmed_up_lr(max_lr=lr, warmup=warmup, global_step=global_step)
    train_op, gnorm = utils.make_optimizer(
      loss=loss,
      trainable_variables=trainable_variables,
      global_step=global_step,
      lr=warmed_up_lr,
      optimizer=optimizer,
      grad_clip=grad_clip / float(num_tpu_replicas()),
      tpu=True
    )

    # ema
    ema, ema_op = make_ema(global_step=global_step, ema_decay=ema_decay, trainable_variables=trainable_variables)
    with tf.control_dependencies([train_op]):
      train_op = tf.group(ema_op)

    # summary
    tpu_summary = TpuSummaries(model_dir, save_summary_steps=100)
    tpu_summary.scalar('train/loss', loss)
    tpu_summary.scalar('train/gnorm', gnorm)
    tpu_summary.scalar('train/pnorm', utils.rms(trainable_variables))
    tpu_summary.scalar('train/lr', warmed_up_lr)
    return tf.estimator.tpu.TPUEstimatorSpec(
      mode=mode, host_call=tpu_summary.get_host_call(), loss=loss, train_op=train_op)

  # Set up Estimator and train
  print("warm_start_from:", warm_start_from)