如何在模型中使用TF指标？

Question

我正在尝试编译一个简单的模型，并且以下工作正常：

model.compile(optimizer=tf.train.AdamOptimizer(0.001),
               loss='categorical_crossentropy',
                metrics=['accuracy'])

但是，我想使用以下指标：tf.metrics.precision_at_k，我不确定如何使它工作。它需要两个参数，这引起了一个问题，因为我不知道如何将预测作为参数传递。我尝试过

model.compile(optimizer=tf.train.AdamOptimizer(0.001),
               loss='categorical_crossentropy',
                metrics=[tf.metrics.precision_at_k])

某些变化但没有帮助。

最低工作代码：

import tensorflow as tf
from tensorflow.keras import layers
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3,3), padding='same', activation=tf.nn.relu, input_shape=(5, 5, 1)),
    tf.keras.layers.MaxPooling2D((2, 2), strides=2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation=tf.nn.relu),
    tf.keras.layers.Dense(2,  activation=tf.nn.softmax)
])