如何在TensorFlow推理会话前批量处理视频流多帧以提升FPS

Great question! You're on the right track—boosting batch size is a common and effective way to get more out of your GPU when it's underutilized while memory is maxed out. Let's walk through how to implement batch processing for your TensorFlow Object Detection API pipeline with OpenCV video streams, tailored to fit typical project structures like yours.

First, start with small batch sizes to avoid hitting memory limits (since you mentioned memory is already full). For your hardware:

• GTX 1050 Laptop: Try batch_size=2 or 4 first, then scale up if memory allows
• Jetson Tx2: Stick to batch_size=2 initially, as embedded GPUs have tighter memory constraints

Test each batch size to find the sweet spot between GPU utilization, FPS, and latency (don't go so big that real-time responsiveness suffers).

Instead of processing one frame at a time, cache frames until you hit your target batch size, then run preprocessing on the entire batch. Here's how to modify your existing loop:

import cv2
import numpy as np

# Configuration
batch_size = 4
input_width, input_height = 640, 640  # Match your model's input size
category_index = # Your category index mapping (from your project)

def preprocess_frame(frame):
    # Reuse your existing preprocessing logic here
    resized = cv2.resize(frame, (input_width, input_height))
    # Convert BGR (OpenCV default) to RGB (model expected format)
    rgb_frame = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
    # Normalize if your model expects it (e.g., /255.0)
    normalized = rgb_frame / 255.0
    return normalized

# Initialize video capture
cap = cv2.VideoCapture(0)  # Replace with your video source
frame_buffer = []

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break
    
    # Store original frame for post-processing (drawing boxes)
    frame_buffer.append(frame.copy())
    
    # Process batch when buffer is full
    if len(frame_buffer) == batch_size:
        # Preprocess all frames in the batch
        batch_input = np.array([preprocess_frame(f) for f in frame_buffer])
        # Shape becomes (batch_size, height, width, 3) - model-compatible

Feed the batch tensor into your session instead of a single frame. The Object Detection API's pre-trained models natively support batch inputs:

# Run batch inference
        feed_dict = {image_tensor: batch_input}
        batch_detections = sess.run(
            [detection_boxes, detection_scores, detection_classes, num_detections],
            feed_dict=feed_dict
        )

Loop through each set of detections in the batch and draw results on the original frames:

# Unpack batch results
        boxes_batch, scores_batch, classes_batch, num_detections_batch = batch_detections
        
        # Process each frame in the batch
        for idx in range(batch_size):
            original_frame = frame_buffer[idx]
            boxes = boxes_batch[idx]
            scores = scores_batch[idx]
            classes = classes_batch[idx]
            num_detections = int(num_detections_batch[idx])
            
            # Reuse your existing draw_boxes function here
            output_frame = draw_boxes(
                original_frame, 
                boxes, 
                scores, 
                classes, 
                num_detections, 
                category_index
            )
            
            # Display or save the frame
            cv2.imshow('Real-Time Detection', output_frame)
        
        # Clear buffer for next batch
        frame_buffer = []
    
    # Handle exit key
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

• Use tf.data.Dataset for Optimized Batching: For even better CPU/GPU synchronization, replace manual frame buffering with TensorFlow's Dataset API. It handles prefetching and batch optimization automatically:

  def frame_generator():
      while cap.isOpened():
          ret, frame = cap.read()
          if not ret:
              break
          yield frame
  
  dataset = tf.data.Dataset.from_generator(
      frame_generator,
      output_types=tf.uint8,
      output_shapes=(None, None, 3)
  )
  dataset = dataset.map(lambda x: preprocess_frame(x)).batch(batch_size).prefetch(tf.data.AUTOTUNE)
  
  # Inference loop
  for batch_input in dataset:
      batch_detections = sess.run(..., feed_dict={image_tensor: batch_input})
      # Post-processing logic here
  

• Asynchronous Frame Reading: If frame capture is slow, use a separate thread to read frames into a queue, so your main thread never waits for new frames to fill the batch.
• Monitor GPU Usage: Use nvidia-smi (for NVIDIA GPUs) or TensorBoard to track utilization after adjusting batch size—stop scaling if memory usage exceeds available capacity.

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