使用TFLite Model Maker训练COCO人物检测模型遇KeyError: 'difficult'

尝试仅使用COCO数据集的人物类别训练TFLite模型，通过Fiftyone处理数据集并导出为VOC格式，执行TFLite Model Maker训练脚本时出现KeyError: 'difficult'错误。

错误信息 #

root@85ac26b47f92:/external# python demofie.py
2022-11-01 21:02:01.059188: E tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: UNKNOWN ERROR (34)
2022-11-01 21:02:01.059234: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:169] retrieving CUDA diagnostic information for host: 85ac26b47f92
2022-11-01 21:02:01.059242: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:176] hostname: 85ac26b47f92
2022-11-01 21:02:01.059324: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:200] libcuda reported version is: NOT_FOUND: was unable to find libcuda.so DSO loaded into this program
2022-11-01 21:02:01.059381: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:204] kernel reported version is: 470.141.3
2022-11-01 21:02:01.059821: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
Traceback (most recent call last):
  File "demofie.py", line 20, in <module>
    train_data = object_detector.DataLoader.from_pascal_voc(images_dir='/external/train/data',annotations_dir='/external/train/labels', label_map=['person'],ignore_difficult_instances= False,num_shards = 100)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_examples/lite/model_maker/core/data_util/object_detector_dataloader.py", line 217, in from_pascal_voc
    cache_writer.write_files(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_examples/lite/model_maker/core/data_util/object_detector_dataloader_util.py", line 252, in write_files
    tf_example = create_pascal_tfrecord.dict_to_tf_example(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow_examples/lite/model_maker/third_party/efficientdet/dataset/create_pascal_tfrecord.py", line 162, in dict_to_tf_example
    if obj['difficult'] == 'Unspecified':
KeyError: 'difficult'

训练代码 #

import numpy as np
import os

from tflite_model_maker.config import QuantizationConfig
from tflite_model_maker.config import ExportFormat
from tflite_model_maker import model_spec
from tflite_model_maker import object_detector

import tensorflow as tf
assert tf.__version__.startswith('2')

tf.get_logger().setLevel('ERROR')
from absl import logging
logging.set_verbosity(logging.ERROR)


spec = model_spec.get('efficientdet_lite1')


train_data = object_detector.DataLoader.from_pascal_voc(images_dir='/external/train/data',annotations_dir='/external/train/labels', label_map=['person'],ignore_difficult_instances= False,num_shards = 100)

validation_data = object_detector.DataLoader.from_pascal_voc(images_dir='/external/val/data',annotations_dir='/external/val/labels',label_map= ['person'],ignore_difficult_instances= False,num_shards = 100)

test_data = object_detector.DataLoader.from_pascal_voc(images_dir='/external/test/data',annotations_dir='/external/test/labels',label_map= ['person'],ignore_difficult_instances= False,num_shards = 100)


model = object_detector.create(train_data, model_spec=spec, batch_size=8,epochs=2000, train_whole_model=True, validation_data=validation_data)

model.evaluate(test_data)

model.export(export_dir='/external/')

数据集生成代码 #

import fiftyone.zoo as foz
import fiftyone as fo
from fiftyone import ViewField as F


cocodataset_test = foz.load_zoo_dataset(
    "coco-2017",
    splits="test",
    label_types=["detections"],
    classes=["person"],
    only_matching=True,
    # max_samples=50,
)
cocodataset_validation = foz.load_zoo_dataset(
    "coco-2017",
    splits="validation",
    label_types=["detections"],
    classes=["person"],
    only_matching=True,
    # max_samples=50
)

cocodataset_train = foz.load_zoo_dataset(
    "coco-2017",
    splits="train",
    label_types=["detections"],
    classes=["person"],
    only_matching=True,
    # max_samples=50,
)


cocodataset_validation.export(
    '/external/val',
    fo.types.VOCDetectionDataset,
)
cocodataset_train.export(
    '/external/train/',
    fo.types.VOCDetectionDataset,
)
cocodataset_test.export(
    '/external/test/',
    fo.types.VOCDetectionDataset,
)

解决方案 #

问题根源 #

Fiftyone导出VOC格式数据集时，默认不为标注添加difficult字段，但TFLite Model Maker的VOC数据加载逻辑会强制读取该字段，导致KeyError。

修复方法（优先推荐） #

在Fiftyone导出数据集前，遍历所有样本的标注，添加difficult字段（VOC格式中0表示非困难样本）：

import fiftyone.zoo as foz
import fiftyone as fo
from fiftyone import ViewField as F


cocodataset_test = foz.load_zoo_dataset(
    "coco-2017",
    splits="test",
    label_types=["detections"],
    classes=["person"],
    only_matching=True,
    # max_samples=50,
)
cocodataset_validation = foz.load_zoo_dataset(
    "coco-2017",
    splits="validation",
    label_types=["detections"],
    classes=["person"],
    only_matching=True,
    # max_samples=50
)

cocodataset_train = foz.load_zoo_dataset(
    "coco-2017",
    splits="train",
    label_types=["detections"],
    classes=["person"],
    only_matching=True,
    # max_samples=50,
)

# 给所有标注添加difficult字段，默认设为0
def add_difficult_field(dataset):
    for sample in dataset:
        for det in sample.detections.detections:
            det["difficult"] = 0
    dataset.save()

add_difficult_field(cocodataset_train)
add_difficult_field(cocodataset_validation)
add_difficult_field(cocodataset_test)

# 执行导出
cocodataset_validation.export(
    '/external/val',
    fo.types.VOCDetectionDataset,
)
cocodataset_train.export(
    '/external/train/',
    fo.types.VOCDetectionDataset,
)
cocodataset_test.export(
    '/external/test/',
    fo.types.VOCDetectionDataset,
)

备选修复方法（修改源码，不推荐） #

找到报错的文件/usr/local/lib/python3.8/dist-packages/tensorflow_examples/lite/model_maker/third_party/efficientdet/dataset/create_pascal_tfrecord.py，将第162行的：

if obj['difficult'] == 'Unspecified':

修改为：

if obj.get('difficult', '0') == 'Unspecified':

这种方式直接给difficult字段设置默认值，但修改依赖包源码会导致后续版本更新失效，仅作为临时方案。

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