为什么TensorFlow-gpu运行YOLOv3目标检测时仅能成功预测第一张图像？

我测试目标检测的示例代码时，出现只有第一张图像可以被成功检测的异常，异常效果如下：

运行环境 #

• TensorFlow-GPU 2.6
• CUDNN 7.6.5
• CUDA Version: 11.4
• Nvidia GeForce RTX 2080 with Max-Q Design

复现代码 #

#! /usr/bin/env python
# coding=utf-8

import cv2
import time
import numpy as np
import core.utils as utils
import tensorflow as tf
from core.yolov3 import YOLOv3, decode


video_path      = "./docs/road.mp4"
num_classes     = 80
input_size      = 416

input_layer  = tf.keras.layers.Input([input_size, input_size, 3])
feature_maps = YOLOv3(input_layer)

bbox_tensors = []
for i, fm in enumerate(feature_maps):
    bbox_tensor = decode(fm, i)
    bbox_tensors.append(bbox_tensor)

model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, "./yolov3.weights")
model.summary()
vid = cv2.VideoCapture(video_path)

Nonstop = True
num = 0
pngfiles = []
while Nonstop:
    num += 1
    return_value, frame = vid.read()
    if return_value:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        framename = "./tmp/image_"+str(num)+".png"
        pngfiles.append(framename)
        cv2.imwrite(framename, frame)
    else:
        print("END!")
        Nonstop = False

num = 0
for framename in pngfiles:
    num += 1
    frame = cv2.imread(framename)
    frame_size = frame.shape[:2]
    image_data = utils.image_preporcess(np.copy(frame), [input_size, input_size])
    image_data = image_data[np.newaxis, ...].astype(np.float32)

    prev_time = time.time()
    #pred_bbox = model.predict_on_batch(image_data)
    pred_bbox = model.predict(image_data, verbose=1)
    curr_time = time.time()
    exec_time = curr_time - prev_time

    pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
    pred_bbox = tf.concat(pred_bbox, axis=0)
    bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size, 0.3)
    bboxes = utils.nms(bboxes, 0.45, method='nms')
    image = utils.draw_bbox(frame, bboxes)

    result = np.asarray(image)
    info = "time: %.2f ms" %(1000*exec_time)
    cv2.putText(result, text=info, org=(50, 70), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                fontScale=1, color=(255, 0, 0), thickness=2)
    cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
    result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite("./tmp/image_predicted_"+str(num)+".png", result)
    cv2.imshow("result", result)
    if cv2.waitKey(1) & 0xFF == ord('q'): break

1. 修复颜色空间不匹配问题：第一阶段存储视频帧时，已经将帧从BGR转为RGB后写入本地，第二阶段用cv2.imread读取图像时，返回的默认是BGR格式，导致输入模型的图像通道顺序反转，特征不匹配。在第二阶段读取图像后新增一行frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)即可解决。
2. 优化推理调用方式：不要在循环中反复调用model.predict，该接口是为批量推理设计的，单张循环调用会导致计算图不断累积，出现推理异常。可以替换为直接调用模型pred_bbox = model(image_data, training=False)，推理效率更高也不会出现图累积问题。
3. 替换循环内的TensorFlow运算：将循环内pred_bbox的reshape、concat操作改为numpy实现，或者整合到模型输出层定义中，避免每次循环都向默认计算图追加节点导致异常。
4. 应急清理方案：如果暂时不想修改核心逻辑，可以在每次循环末尾新增tf.keras.backend.clear_session()清理会话缓存，即可恢复后续帧的推理效果，但会降低整体运行速度。

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