使用SSD与MobileNetV2训练目标检测模型时形状不匹配问题求助

问题报错 #

训练时出现形状不匹配错误：

ValueError: Arguments `target` and `output` must have the same shape. Received: target.shape=(None, 1, 2), output.shape=(None, 49, 2)

问题根源 #

SSD是基于**锚框（Anchor Boxes）**的检测算法，你的模型输出对应7x7=49个锚框的分类和回归结果（输入224x224时，MobileNetV2输出7x7特征图），但数据解析只生成了1个真实框的标签，形状完全不匹配。必须让每个锚框都对应一个标签（标记是否为正样本、对应的类别和框偏移）。

解决方案 #

1. 修正数据解析代码，生成锚框匹配的标签 #

修改TFRecord解析逻辑，为每个样本生成49个锚框对应的标签：

import tensorflow as tf

def generate_anchors():
    # 生成7x7网格的基础锚框（每个网格1个，后续可扩展为多个比例）
    grid_size = 7
    anchor_size = 32  # 224/7=32，单元格尺寸
    anchors = []
    for y in range(grid_size):
        for x in range(grid_size):
            # 锚框中心归一化坐标
            cx = (x + 0.5) / grid_size
            cy = (y + 0.5) / grid_size
            # 锚框宽高归一化
            w = anchor_size / 224
            h = anchor_size / 224
            # 转换为xmin, ymin, xmax, ymax格式
            xmin = cx - w/2
            ymin = cy - h/2
            xmax = cx + w/2
            ymax = cy + h/2
            anchors.append([xmin, ymin, xmax, ymax])
    return tf.convert_to_tensor(anchors, dtype=tf.float32)

def compute_iou(box1, boxes2):
    # 计算单个真实框与所有锚框的IOU
    x1 = tf.maximum(box1[0], boxes2[:, 0])
    y1 = tf.maximum(box1[1], boxes2[:, 1])
    x2 = tf.minimum(box1[2], boxes2[:, 2])
    y2 = tf.minimum(box1[3], boxes2[:, 3])
    
    area_inter = tf.maximum(0.0, x2 - x1) * tf.maximum(0.0, y2 - y1)
    area_box1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area_box2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
    area_union = area_box1 + area_box2 - area_inter
    
    return area_inter / area_union

def parse_tf_example(example_proto):
    feature_description = {
        'image/height': tf.io.FixedLenFeature([], tf.int64),
        'image/width': tf.io.FixedLenFeature([], tf.int64),
        'image/encoded': tf.io.FixedLenFeature([], tf.string),
        'image/format': tf.io.FixedLenFeature([], tf.string),
        'image/object/bbox/xmin': tf.io.FixedLenFeature([], tf.float32),
        'image/object/bbox/ymin': tf.io.FixedLenFeature([], tf.float32),
        'image/object/bbox/xmax': tf.io.FixedLenFeature([], tf.float32),
        'image/object/bbox/ymax': tf.io.FixedLenFeature([], tf.float32),
        'image/object/class/label': tf.io.FixedLenFeature([], tf.int64)
    }
    parsed_features = tf.io.parse_single_example(example_proto, feature_description)
    
    # 图像预处理
    image = tf.image.decode_jpeg(parsed_features['image/encoded'], channels=3)
    image = tf.image.resize(image, [224, 224])
    image = tf.cast(image, tf.float32) / 255.0
    
    # 提取真实框和标签
    true_bbox = tf.convert_to_tensor([
        parsed_features['image/object/bbox/xmin'],
        parsed_features['image/object/bbox/ymin'],
        parsed_features['image/object/bbox/xmax'],
        parsed_features['image/object/bbox/ymax']
    ], dtype=tf.float32)
    true_label = tf.one_hot(parsed_features['image/object/class/label'], depth=2)
    
    # 生成锚框并匹配
    anchors = generate_anchors()  # shape (49,4)
    ious = compute_iou(true_bbox, anchors)
    max_iou_idx = tf.argmax(ious)  # 找到与真实框IOU最大的锚框作为正样本
    
    # 初始化分类标签：负样本全为0，正样本设为真实标签
    class_labels = tf.zeros((49, 2), dtype=tf.float32)
    class_labels = tf.tensor_scatter_nd_update(class_labels, [[max_iou_idx]], [true_label])
    
    # 生成边界框回归目标：仅正样本计算偏移，负样本为0
    anchor = anchors[max_iou_idx]
    anchor_w = anchor[2] - anchor[0]
    anchor_h = anchor[3] - anchor[1]
    # SSD标准回归偏移计算
    dx = (true_bbox[0] - anchor[0]) / anchor_w
    dy = (true_bbox[1] - anchor[1]) / anchor_h
    dw = tf.math.log((true_bbox[2] - true_bbox[0]) / anchor_w)
    dh = tf.math.log((true_bbox[3] - true_bbox[1]) / anchor_h)
    
    bbox_targets = tf.zeros((49, 4), dtype=tf.float32)
    bbox_targets = tf.tensor_scatter_nd_update(bbox_targets, [[max_iou_idx]], [[dx, dy, dw, dh]])
    
    # 合并分类和回归标签，形状为(49,6)，与模型输出匹配
    combined_label = tf.concat([class_labels, bbox_targets], axis=-1)
    return image, combined_label

# 加载数据集代码不变
def load_tfrecords(tfrecord_path):
    raw_dataset = tf.data.TFRecordDataset(tfrecord_path)
    parsed_dataset = raw_dataset.map(parse_tf_example, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    return parsed_dataset

train_dataset = load_tfrecords('train.tfrecord')
val_dataset = load_tfrecords('val.tfrecord')
train_dataset = train_dataset.batch(32).prefetch(tf.data.experimental.AUTOTUNE)
val_dataset = val_dataset.batch(32).prefetch(tf.data.experimental.AUTOTUNE)

2. 修正损失函数，仅计算正样本的边界框损失 #

负样本的边界框回归没有意义，需要过滤掉这部分损失：

from tensorflow.keras.losses import BinaryCrossentropy, MeanSquaredError

def ssd_loss(y_true, y_pred):
    num_classes = 2
    num_bboxes = 4
    
    class_true = y_true[..., :num_classes]
    bbox_true = y_true[..., num_classes:]
    class_pred = y_pred[..., :num_classes]
    bbox_pred = y_pred[..., num_classes:]
    
    # 标记正样本：分类标签非全0的锚框
    positive_mask = tf.reduce_any(class_true > 0.5, axis=-1)
    positive_mask = tf.expand_dims(positive_mask, axis=-1)  # 扩展维度以便广播
    
    # 分类损失：所有样本参与计算
    class_loss = BinaryCrossentropy()(class_true, class_pred)
    
    # 边界框损失：仅正样本参与计算
    bbox_loss = MeanSquaredError()(bbox_true * positive_mask, bbox_pred * positive_mask)
    
    # 给边界框损失加权重，提升回归优先级
    return class_loss + 5.0 * bbox_loss

# 编译和训练代码不变
model.compile(optimizer=Adam(learning_rate=1e-4), loss=ssd_loss)

callbacks = [
    tf.keras.callbacks.ModelCheckpoint('best_model.keras', save_best_only=True, monitor='val_loss'),
    tf.keras.callbacks.EarlyStopping(patience=5, monitor='val_loss')
]

history = model.fit(
    train_dataset,
    validation_data=val_dataset,
    epochs=30,
    callbacks=callbacks
)

3. 可选优化：增加多比例锚框 #

SSD通常每个网格设置多个不同比例的锚框（比如3个），提升对不同尺寸目标的检测能力。只需修改generate_anchors函数，为每个网格生成多个锚框即可，同时调整模型输出的锚框数量（比如7x7x3=147个锚框，对应class_output和bbox_output的通道数调整为23和43，Reshape后为(147,2)和(147,4)）。

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