如何在Keras中创建和使用加权指标？

import tensorflow as tf
from tensorflow import keras

class WeightedAccuracy(keras.metrics.Metric):
    def __init__(self, name='weighted_accuracy', **kwargs):
        super(WeightedAccuracy, self).__init__(name=name, **kwargs)
        self.total = self.add_weight(name='total', initializer='zeros')
        self.count = self.add_weight(name='count', initializer='zeros')

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_pred = tf.argmax(y_pred, axis=-1)
        values = tf.cast(tf.equal(y_true, y_pred), tf.float32)
        if sample_weight is not None:
            values *= sample_weight
        self.total.assign_add(tf.reduce_sum(values))
        self.count.assign_add(tf.reduce_sum(sample_weight))

    def result(self):
        return self.total / self.count

    def reset_states(self):
        self.total.assign(0.0)
        self.count.assign(0.0)

# 使用加权指标进行模型编译和训练
model = keras.models.Sequential([...])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=[WeightedAccuracy()])

# 加载数据集并进行训练
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train, sample_weights))
model.fit(train_dataset, epochs=10)