如何基于TFRecord文件训练自定义目标检测模型并适配自建模型

1. 核心文件作用说明 #

• TFRecord：Roboflow导出的打包数据集，包含图像二进制数据和标注信息（边界框坐标、类别ID）
• .pbtxt：类别映射文件，定义类别与ID的对应关系，示例格式：

  item {
    id: 1
    name: 'cat'
  }
  item {
    id: 2
    name: 'dog'
  }
  

  训练时需保证模型输出的类别ID与该文件完全匹配。

2. 适配TFRecord的解析函数 #

TFRecord无法直接输入模型，需先解析为张量格式，以下是适配Roboflow导出格式的解析代码：

import tensorflow as tf
from object_detection.utils import label_map_util

# 加载类别映射
label_map = label_map_util.load_labelmap('your_classes.pbtxt')
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=3, use_display_name=True)
num_classes = len(categories)

def parse_tfrecord_fn(example):
    # 定义TFRecord特征结构（与Roboflow导出格式对齐）
    feature_desc = {
        'image/encoded': tf.io.FixedLenFeature([], tf.string),
        'image/object/bbox/xmin': tf.io.VarLenFeature(tf.float32),
        'image/object/bbox/xmax': tf.io.VarLenFeature(tf.float32),
        'image/object/bbox/ymin': tf.io.VarLenFeature(tf.float32),
        'image/object/bbox/ymax': tf.io.VarLenFeature(tf.float32),
        'image/object/class/label': tf.io.VarLenFeature(tf.int64),
    }
    example = tf.io.parse_single_example(example, feature_desc)
    
    # 解码并预处理图像
    image = tf.io.decode_jpeg(example['image/encoded'], channels=3)
    image = tf.image.resize(image, (224, 224))  # 统一输入尺寸，可按需调整
    image = image / 255.0  # 归一化到0-1区间
    
    # 转换边界框与标签格式
    xmin = tf.sparse.to_dense(example['image/object/bbox/xmin'])
    xmax = tf.sparse.to_dense(example['image/object/bbox/xmax'])
    ymin = tf.sparse.to_dense(example['image/object/bbox/ymin'])
    ymax = tf.sparse.to_dense(example['image/object/bbox/ymax'])
    bboxes = tf.stack([ymin, xmin, ymax, xmax], axis=1)  # 转为(y1,x1,y2,x2)格式
    labels = tf.sparse.to_dense(example['image/object/class/label'])
    
    return image, {'bbox': bboxes, 'label': labels}

3. 构建仅含2-3个隐藏层的轻量模型 #

采用"轻量骨干网络+精简检测头"结构，其中2-3个隐藏层放在检测头部分：

def build_lightweight_detector(input_shape=(224,224,3), num_classes=3):
    # 输入层
    inputs = tf.keras.Input(shape=input_shape)
    
    # 骨干网络：简单卷积特征提取
    x = tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding='same')(inputs)
    x = tf.keras.layers.MaxPooling2D((2,2))(x)
    x = tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding='same')(x)
    x = tf.keras.layers.MaxPooling2D((2,2))(x)
    
    # 2-3个隐藏层（检测头核心）
    x = tf.keras.layers.Conv2D(128, (3,3), activation='relu', padding='same')(x)  # 第1个隐藏层
    x = tf.keras.layers.Conv2D(64, (3,3), activation='relu', padding='same')(x)   # 第2个隐藏层
    # 可选第3个隐藏层：x = tf.keras.layers.Conv2D(32, (3,3), activation='relu', padding='same')(x)
    
    # 检测头输出：边界框回归+类别分类（适配多目标锚框思路）
    bbox_output = tf.keras.layers.Conv2D(4*9, (3,3), activation='linear', padding='same')(x)  # 9个锚框，每个输出4个坐标
    class_output = tf.keras.layers.Conv2D(num_classes*9, (3,3), activation='softmax', padding='same')(x)
    
    return tf.keras.Model(inputs=inputs, outputs={'bbox': bbox_output, 'label': class_output})

4. 数据集管道与训练流程 #

4.1 构建训练/验证数据集 #

# 加载TFRecord文件
train_ds = tf.data.TFRecordDataset('train.tfrecord')
val_ds = tf.data.TFRecordDataset('val.tfrecord')

# 解析并预处理
train_ds = train_ds.map(parse_tfrecord_fn).shuffle(100).batch(8)
val_ds = val_ds.map(parse_tfrecord_fn).batch(8)

4.2 定义损失与训练 #

目标检测损失为边界框回归损失与类别分类损失的加权和：

def detection_loss(y_true, y_pred):
    # 边界框L1损失
    bbox_loss = tf.keras.losses.MeanAbsoluteError()(y_true['bbox'], y_pred['bbox'])
    # 类别交叉熵损失
    label_loss = tf.keras.losses.SparseCategoricalCrossentropy()(y_true['label'], y_pred['label'])
    # 加权总损失
    return 0.5 * bbox_loss + 0.5 * label_loss

# 编译并训练
model = build_lightweight_detector(num_classes=num_classes)
model.compile(optimizer=tf.keras.optimizers.Adam(1e-4), loss=detection_loss)

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=50,
    callbacks=[tf.keras.callbacks.EarlyStopping(patience=5, restore_best_weights=True)]
)

5. 关键注意事项 #

• 若为单目标检测，可简化检测头，无需锚框，直接输出1组边界框和类别
• 若解析TFRecord时报错，可打印example查看字段名，与Roboflow导出格式对齐
• 隐藏层数量与通道数需匹配数据集大小：数据集小则减少参数，避免过拟合

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