ViT模型添加准确率外指标（召回率）时出现报错求助

问题背景 #

基于Keras的ViT图像分类示例代码，使用自定义多分类数据集训练模型，仅使用accuracy指标时训练正常，添加Recall指标后出现训练报错。

报错信息 #

/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52   try:
     53     ctx.ensure_initialized()
---> 54     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                         inputs, attrs, num_outputs)
     56   except core._NotOkStatusException as e:

InvalidArgumentError: Graph execution error:

Detected at node 'assert_greater_equal/Assert/AssertGuard/Assert' defined at (most recent call last):
    File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
      return _run_code(code, main_globals, None,
    File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
      exec(code, run_globals)
    File "/usr/local/lib/python3.8/dist-packages/ipykernel_launcher.py", line 16, in <module>
      app.launch_new_instance()
    File "/usr/local/lib/python3.8/dist-packages/traitlets/config/application.py", line 992, in launch_instance
      app.start()
    File "/usr/local/lib/python3.8/dist-packages/ipykernel/kernelapp.py", line 612, in start
      self.io_loop.start()
    File "/usr/local/lib/python3.8/dist-packages/tornado/platform/asyncio.py", line 149, in start
      self.asyncio_loop.run_forever()
...

(1) INVALID_ARGUMENT:  assertion failed: [predictions must be >= 0] [Condition x >= y did not hold element-wise:] [x (model_3/dense_79/BiasAdd:0) = ] [[251.636795 -233.491394 322.750397...]...] [y (Cast_4/x:0) = ] [0]
     [[{{node assert_greater_equal/Assert/AssertGuard/Assert}}]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_348766]

数据集情况 #

已将标签转为one-hot格式，数组形状：

x_train shape: (6179, 336, 336, 3) - y_train shape: (6179, 9) x_test shape: (2060, 336, 336, 3) - y_test shape: (2060, 9)

修改的代码片段 #

optimizer = tfa.optimizers.AdamW(
    learning_rate=learning_rate, weight_decay=weight_decay
)

model.compile(
        optimizer=optimizer,
        loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
        metrics = [ 'accuracy', tf.keras.metrics.Recall(name = 'recall')]
)

checkpoint_filepath = "/content/drive/MyDrive/DATASET/checkpoints/checkpoint.h5"
checkpoint_callback = keras.callbacks.ModelCheckpoint(
        checkpoint_filepath,
        monitor="val_accuracy",
        save_best_only=True,
        save_weights_only=True,
    )

r = model.fit(
        x=x_train,
        y=y_train,
        batch_size=batch_size,
        epochs=num_epochs,
        validation_split=0.25,
)

model.load_weights(checkpoint_filepath)
_, accuracy = model.evaluate(x_test, y_test)
print(f"Test accuracy: {round(accuracy * 100, 2)}%")

解决方法 #

报错核心原因：Recall指标默认期望输入是0-1之间的概率值，但当前模型输出是未经过Softmax激活的logits（可能出现负数），而损失函数设置了from_logits=True，说明模型最后一层没有激活函数。

提供两种可行解决方案：

方案1：修改Recall指标参数适配logits输入 #

在初始化Recall指标时，指定from_logits=True，同时针对多分类任务设置average参数（可选macro/micro/weighted，根据需求选择），示例代码：

metrics = [
    'accuracy', 
    tf.keras.metrics.Recall(name='recall', from_logits=True, average='macro')
]

注：多分类任务必须指定average，否则会默认使用二分类的binary模式导致报错。

方案2：给模型输出层添加Softmax激活 #

修改模型最后一层，添加softmax激活函数，同时移除损失函数的from_logits=True参数：

• 模型输出层（假设原输出层为Dense(9)）改为：tf.keras.layers.Dense(9, activation='softmax')
• 损失函数改为：tf.keras.losses.CategoricalCrossentropy()

两种方案任选其一即可，推荐方案1，因为直接使用logits训练具备更好的数值稳定性。

内容的提问来源于stack exchange，提问作者gulia