列表输出的Keras自定义精度度量

import tensorflow as tf
from tensorflow import keras

class CustomAccuracy(keras.metrics.Metric):
    def __init__(self, name='custom_accuracy', **kwargs):
        super(CustomAccuracy, self).__init__(name=name, **kwargs)
        self.true_positives = self.add_weight(name='tp', initializer='zeros')
        self.total_samples = self.add_weight(name='ts', initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_pred = tf.argmax(y_pred, axis=-1)
        y_true = tf.cast(y_true, tf.int64)
        values = tf.equal(y_true, y_pred)
        values = tf.cast(values, tf.float32)

        if sample_weight is not None:
            sample_weight = tf.cast(sample_weight, tf.float32)
            values = tf.multiply(values, sample_weight)

        self.true_positives.assign_add(tf.reduce_sum(values))
        self.total_samples.assign_add(tf.cast(tf.size(y_true), tf.float32))

    def result(self):
        return self.true_positives / self.total_samples

    def reset_states(self):
        self.true_positives.assign(0.0)
        self.total_samples.assign(0.0)

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=[CustomAccuracy()])