Как лучше всего создать матрицу путаницы с несколькими метками на TensorFlow?

Question

Я пытаюсь сгенерировать путаницу с 3 метками, которые можно предсказать с помощью функции TensorFlow tf.confusion_matrix().

Моя проблема в том, что способ, которым я это реализовал, очень медленный.

Я делаю это для вызова функции, которая перебирает набор данных и суммирует матрицу путаницы для каждого пакета. Я думаю, проблема в том, что мой тензор tf.confusion_matrix определен в отдельном графе, который требует его оценки с использованием другого сеанса (см. Код ниже).

Я попытался реализовать его в основном графике, но не понял, как сгладить (развернуть) прогнозы в списке.

def do_accuracy_eval(sess, data_set, data_index, data_type, batch_size, self):    
steps_per_epoch = len(data_set) // batch_size
for step in range(steps_per_epoch):
    batch_inputs, batch_labels, batch_info, data_index, word_labels = generate_batch(batch_size, 
                                                                        data_set, 
                                                                        data_index)

    feed_dict = {self._train_inputs : batch_inputs, self._train_labels : batch_labels, 
                 self._info_inputs : batch_info, self._train_labels_NER : word_labels} 

    prediction_NER = sess.run([self._prediction_NER], feed_dict=feed_dict)

    #compute the confusion matrix for NER
    word_labels_unstacked =        word_labels.reshape(word_labels.shape[0]*word_labels.shape[1])
    prediction_NER_unstacked = prediction_NER.reshape(prediction_NER.shape[0]*prediction_NER.shape[1])
    if step == 0:
        NER_confusion_mat = tf.confusion_matrix(word_labels_unstacked, prediction_NER_unstacked, num_classes=3, name='NER_confusion_matrix')
    else:
        NER_confusion_mat += tf.confusion_matrix(word_labels_unstacked, prediction_NER_unstacked, num_classes=3, name='NER_confusion_matrix')

print('confusion matrix NER')
session = tf.Session()
with session.as_default():
    cm = session.run(NER_confusion_mat)
    print(cm)

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

prediction = tf.cast(tf.argmax(logits, axis=1), tf.float32, name='prediction')
t_p = tf.count_nonzero(prediction * tf.cast(train_labels, tf.float32), name='tp')
t_n = tf.count_nonzero((prediction - 1) * (tf.cast(train_labels, tf.float32) - 1), name='tn')
f_p = tf.count_nonzero(prediction * (tf.cast(train_labels, tf.float32) - 1), name='fp')
f_n = tf.count_nonzero((prediction - 1) * tf.cast(train_labels, tf.float32), name='fn')

Спасибо за вашу помощь!