Использование пакетной нормализации с выпадением на этапе тестирования дает Nan

Я пытаюсь добавить пакетную нормализацию в свой CNN и прочитал много сообщений о том, как это сделать, но все же моя реализация дает в результате массив нанов при установке для обучения значения False.

При установке для обучения значения True даже во время тестирования результаты не равны Nan, но они хуже, чем во время обучения, если я тестирую на обучающих изображениях.

Я использовал спад 0,9 и тренировался на 15 000 итераций.

Вот мое построение графика, добавление операций обновления в качестве зависимости, как рекомендовано в tf.layers .batch_normalization documentation, а затем с помощью sessions

extra_update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
  with tf.control_dependencies(extra_update_ops):
    phase_train = tf.placeholder(tf.bool, name='phase_train')

    ###### Other placeholders and variables declarations ######

    # Build a Graph that computes the logits predictions from the inference model.

    loss, eval_prediction = inference(train_data_node, train_labels_node, batch_size, phase_train, dropout_out_keep_prob)

    # Build a Graph that trains the model with one batch of examples and updates the model parameters.

    ###### Should I rather put the dependency here ? ######
    train_op = train(loss, global_step)

    saver = tf.train.Saver(tf.global_variables())

    with tf.Session() as sess:
          init = tf.global_variables_initializer()
          sess.run(init)

          # Start the queue runners.
          coord = tf.train.Coordinator()
          threads = tf.train.start_queue_runners(sess=sess, coord=coord)

          for step in range(startstep, startstep + max_steps):
            feed_dict = fill_feed_dict(train_labels_node, train_data_node, dropout_out_keep_prob, phase_train, batch_size, phase_train_val=True,drop_out_keep_prob_val=1.)
            _, loss_value = sess.run([train_op, loss], feed_dict=feed_dict)

вот мой вызов функции batch_norm:

def batch_norm_layer(inputT, is_training, scope):
    return tf.layers.batch_normalization(inputT, training=is_training, center=False, reuse=None, momentum=0.9)

а теперь вот как я восстанавливаю модель для тестирования:

phase_train = tf.placeholder(tf.bool, name='phase_train')

###### Other placeholder definitions ######

loss, logits = inference(test_data_node, test_labels_node, batch_size, phase_train, drop_out_keep_prob)
pred = tf.argmax(logits, dimension=3)

saver = tf.train.Saver()

with tf.Session() as sess:
  saver.restore(sess, test_ckpt)

  threads = tf.train.start_queue_runners(sess=sess)

  feed_dict = fill_feed_dict(test_labels_node, test_data_node, drop_out_keep_prob, phase_train, batch_size=1, phase_train_val=False, drop_out_keep_prob_val=1.)

  pred_loss, dense_prediction, predicted_image = sess.run([loss, logits, pred], feed_dict=feed_dict)

Здесь density_prediction дает массив Nans и, таким образом, predicted_image имеет значение 0. Есть ли ошибка в моей конструкции? Как я могу это исправить / диагностировать?

Любая помощь будет приветствоваться, я прочитал много руководств по использованию "ручной" пакетной нормы, но я не смог найти хорошо управляемый учебник о том, как использовать официальную пакетную норму, думаю, это потому, что это слишком очевидно, но это не для меня !