将YOLOv8替换InceptionV3集成至Transformer图像字幕模型求助

我有一个基于Transformer的图像字幕生成系统，原本使用InceptionV3提取图像特征。现在需要将YOLOv8模型替换InceptionV3，把YOLOv8检测到的目标及提取的特征传入Transformer生成字幕。

原有InceptionV3特征提取代码 #

def CNN_Encoder_Incep():
    inception_v3 = tf.keras.applications.InceptionV3(
        include_top=False,
        weights='imagenet'
    )
    inception_v3.trainable = False

    output = inception_v3.output
    output = tf.keras.layers.Reshape(
        (-1, output.shape[-1]))(output)

    cnn_model = tf.keras.models.Model(inception_v3.input, output)
    return cnn_model

class ImageCaptioningModel(tf.keras.Model):

    def __init__(self, cnn_model, encoder, decoder, image_aug=None):
        super().__init__()
        self.cnn_model = cnn_model
        self.encoder = encoder
        self.decoder = decoder
        self.image_aug = image_aug
        self.loss_tracker = tf.keras.metrics.Mean(name="loss")
        self.acc_tracker = tf.keras.metrics.Mean(name="accuracy")


    def calculate_loss(self, y_true, y_pred, mask):
        loss = self.loss(y_true, y_pred)
        mask = tf.cast(mask, dtype=loss.dtype)
        loss *= mask
        return tf.reduce_sum(loss) / tf.reduce_sum(mask)


    def calculate_accuracy(self, y_true, y_pred, mask):
        accuracy = tf.equal(y_true, tf.argmax(y_pred, axis=2))
        accuracy = tf.math.logical_and(mask, accuracy)
        accuracy = tf.cast(accuracy, dtype=tf.float32)
        mask = tf.cast(mask, dtype=tf.float32)
        return tf.reduce_sum(accuracy) / tf.reduce_sum(mask)


    def compute_loss_and_acc(self, img_embed, captions, training=True):
        encoder_output = self.encoder(img_embed, training=True)
        y_input = captions[:, :-1]
        y_true = captions[:, 1:]
        mask = (y_true != 0)
        y_pred = self.decoder(
            y_input, encoder_output, training=True, mask=mask
        )
        loss = self.calculate_loss(y_true, y_pred, mask)
        acc = self.calculate_accuracy(y_true, y_pred, mask)
        return loss, acc


    def train_step(self, batch):
        imgs, captions = batch

        if self.image_aug:
            imgs = self.image_aug(imgs)

        img_embed = self.cnn_model(imgs)

        with tf.GradientTape() as tape:
            loss, acc = self.compute_loss_and_acc(
                img_embed, captions
            )

        train_vars = (
            self.encoder.trainable_variables + self.decoder.trainable_variables
        )
        grads = tape.gradient(loss, train_vars)
        self.optimizer.apply_gradients(zip(grads, train_vars))
        self.loss_tracker.update_state(loss)
        self.acc_tracker.update_state(acc)

        return {"loss": self.loss_tracker.result(), "acc": self.acc_tracker.result()}


    def test_step(self, batch):
        imgs, captions = batch

        img_embed = self.cnn_model(imgs)

        loss, acc = self.compute_loss_and_acc(
            img_embed, captions, training=False
        )

        self.loss_tracker.update_state(loss)
        self.acc_tracker.update_state(acc)

        return {"loss": self.loss_tracker.result(), "acc": self.acc_tracker.result()}

    @property
    def metrics(self):
        return [self.loss_tracker, self.acc_tracker]

cnn_model = CNN_Encoder_Incep()
caption_model = ImageCaptioningModel(
    cnn_model=cnn_model, encoder=encoder, decoder=decoder, image_aug=image_augmentation,
)

尝试的YOLOv8替换代码（报错） #

def CNN_Encoder():
    yolov8_model = tf.keras.models.load_model('./content/yolov8n_objdet_oidv7_640x640.pt')
    yolov8_model.trainable = False
    output = yolov8_model.output
    output = tf.keras.layers.Reshape((-1, output.shape[-1]))(output)
    cnn_model = tf.keras.models.Model(yolov8_model.input, output)
    cnn_model_onnx = cnn_model.export(format='onnx')
    return cnn_model

问题分析与解决方案 #

核心错误原因 #

1. 模型加载方式错误：tf.keras.models.load_model无法直接加载YOLOv8的.pt权重文件，这是Ultralytics专属格式，并非Keras SavedModel格式。
2. 输出结构不匹配：YOLOv8原始输出是检测框、置信度和分类结果，并非图像特征图，直接Reshape会导致特征维度与Transformer编码器输入不兼容。
3. 冗余操作干扰：特征提取阶段不需要导出ONNX，该操作会打断模型返回流程。

正确的YOLOv8特征提取实现 #

使用Ultralytics官方API加载模型，提取Backbone层的特征输出，并处理成与原InceptionV3兼容的维度：

from ultralytics import YOLO
import tensorflow as tf

def CNN_Encoder_YOLOv8():
    # 加载预训练YOLOv8模型
    yolov8_model = YOLO('./content/yolov8n_objdet_oidv7_640x640.pt')
    
    # 获取YOLOv8 Backbone最后一层的输出作为图像特征
    backbone_output = yolov8_model.model.model[0][-1].output
    
    # 构建Keras特征提取模型
    cnn_model = tf.keras.models.Model(inputs=yolov8_model.model.input, outputs=backbone_output)
    cnn_model.trainable = False
    
    # 调整特征形状为 (batch_size, 特征数, 特征维度)，匹配原InceptionV3输出格式
    # YOLOv8n Backbone输出为 (batch_size, 20, 20, 1280)，Reshape后为 (batch_size, 400, 1280)
    output = tf.keras.layers.Reshape((-1, backbone_output.shape[-1]))(cnn_model.output)
    
    # 可选：添加Dense层将特征维度转换为原InceptionV3的2048，避免修改Transformer编码器参数
    # output = tf.keras.layers.Dense(2048, activation='relu')(output)
    
    cnn_model = tf.keras.models.Model(cnn_model.input, output)
    return cnn_model

适配Transformer的关键注意事项 #

1. 特征维度对齐：若不想修改Transformer编码器的输入维度，可在Reshape后添加Dense层将YOLOv8的1280维特征转换为2048维。
2. 图像预处理适配：YOLOv8的输入要求与InceptionV3不同，需替换图像增强/预处理函数：

   def yolov8_preprocess(imgs):
       # 调整图像尺寸到YOLOv8默认的640x640
       imgs = tf.image.resize(imgs, (640, 640))
       # 归一化到0-1区间（YOLOv8默认预处理逻辑）
       imgs = imgs / 255.0
       return imgs
   

3. 版本兼容性：确保ultralytics与TensorFlow版本兼容，建议使用最新版ultralytics。

替换后的完整调用 #

cnn_model = CNN_Encoder_YOLOv8()
caption_model = ImageCaptioningModel(
    cnn_model=cnn_model, encoder=encoder, decoder=decoder, image_aug=yolov8_preprocess,
)

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