我正在训练一个分类模型，我决定从类别交叉熵损失函数切换到稀疏类别交叉熵，这样可以使用更少的内存，加快训练速度。我的训练计算精度和召回率指标。
然而，当我切换到稀疏交叉熵时，精度度量开始失败。问题是SparseCategoricalCrossentropy期望真标签是标量，而预测标签是大小为“类数”的矢量，精度度量引发了“形状不匹配”类型的异常。
一个最小的例子来说明这一点（相同的模型在没有精确度得分的情况下工作，并且在增加了精确度得分计算的第二次训练期间失败）：

import numpy as np
import tensorflow as tf

x = np.arange(0, 20)
y = np.zeros_like(x)
for i in range(len(x)):
    if x[i] % 2 == 0:
        y[i] = 0  # Even number
    else:
        y[i] = 1  # Odd number
n_classes = len(np.unique(y))

model = tf.keras.Sequential(
    [
        tf.keras.layers.Dense(10, input_shape=(1,)),
        tf.keras.layers.Dense(n_classes, activation="softmax"),
    ]
)

print("Train without precision metric")
model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
)
model.fit(x, y, epochs=2)

print("Train with precision metric")
model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=[tf.keras.metrics.Precision()],
)
model.fit(x, y, epochs=2)

输出为

Metal device set to: Apple M1 Pro
2022-09-20 18:47:20.254419: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.
2022-09-20 18:47:20.254522: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)
2022-09-20 18:47:20.324585: W tensorflow/core/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
Train without precision metric
Epoch 1/2
2022-09-20 18:47:20.441786: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:113] Plugin optimizer for device_type GPU is enabled.

1/1 [==============================] - ETA: 0s - loss: 5.9380
1/1 [==============================] - 0s 205ms/step - loss: 5.9380
Epoch 2/2

1/1 [==============================] - ETA: 0s - loss: 5.8844
1/1 [==============================] - 0s 4ms/step - loss: 5.8844
Train with precision metric
Epoch 1/2

systemMemory: 16.00 GB
maxCacheSize: 5.33 GB

Traceback (most recent call last):
  File "/Users/dima/dev/learn/datascience/test-sparse-precision.py", line 35, in <module>
    model.fit(x, y, epochs=2)
  File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/var/folders/_0/2yc8qfs11xq2vykxzkkngq4m0000gn/T/__autograph_generated_filedw4nh8_p.py", line 15, in tf__train_function
    retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
ValueError: in user code:

    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/engine/training.py", line 1051, in train_function  *
        return step_function(self, iterator)
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/engine/training.py", line 1040, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/engine/training.py", line 1030, in run_step  **
        outputs = model.train_step(data)
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/engine/training.py", line 894, in train_step
        return self.compute_metrics(x, y, y_pred, sample_weight)
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/engine/training.py", line 987, in compute_metrics
        self.compiled_metrics.update_state(y, y_pred, sample_weight)
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/engine/compile_utils.py", line 501, in update_state
        metric_obj.update_state(y_t, y_p, sample_weight=mask)
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/utils/metrics_utils.py", line 70, in decorated
        update_op = update_state_fn(*args, **kwargs)
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/metrics/base_metric.py", line 140, in update_state_fn
        return ag_update_state(*args, **kwargs)
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/metrics/metrics.py", line 818, in update_state  **
        return metrics_utils.update_confusion_matrix_variables(
    File "/Users/dima/sw/mambaforge/envs/data-science/lib/python3.10/site-packages/keras/utils/metrics_utils.py", line 619, in update_confusion_matrix_variables
        y_pred.shape.assert_is_compatible_with(y_true.shape)

    ValueError: Shapes (None, 2) and (None, 1) are incompatible

它发生在两种不同的环境中：苹果公司的Tensorflow 2.9.2和Ubuntu上的Tensorflow 2.8.0。
除了编写我自己的度量类之外，还有人知道如何处理这个问题吗？