如何在Tesla K80上将目标检测预测速度提升20倍？

Alright, let's tackle your goal of getting a 20x speedup on your Tesla K80 for object detection inference. First, let's break down the low-hanging fruit and then move to more impactful optimizations tailored to your setup and code.

Problem Statement #

I'm running object detection inference on a Tesla K80, and the FPS is way too low—currently taking ~1.8 seconds per image (including visualization and saving). I only need the output_dict and don't care about visualization or saving images. I need to speed this up by at least 20x.

Here's my current code:

def load_detection_graph(PATH_TO_CKPT):
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    return detection_graph

def load_image_into_numpy_array(image):
    ''' convert image to numpy arrays '''
    (im_width, im_height) = image.size
    return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)

def run_inference_for_single_image(image, graph, filename):
    with graph.as_default():
        with tf.Session() as sess:
            # Get handles to input and output tensors
            ops = tf.get_default_graph().get_operations()
            all_tensor_names = {output.name for op in ops for output in op.outputs}
            tensor_dict = {}
            for key in [
                'num_detections', 'detection_boxes', 'detection_scores',
                'detection_classes', 'detection_masks'
            ]:
                tensor_name = key + ':0'
                if tensor_name in all_tensor_names:
                    tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
                        tensor_name)
            if 'detection_masks' in tensor_dict:
                # The following processing is only for single image
                detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
                detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
                # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
                real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
                detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
                detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
                detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
                    detection_masks, detection_boxes, image.shape[0], image.shape[1])
                detection_masks_reframed = tf.cast(
                    tf.greater(detection_masks_reframed, 0.5), tf.uint8)
                # Follow the convention by adding back the batch dimension
                tensor_dict['detection_masks'] = tf.expand_dims(
                    detection_masks_reframed, 0)
            image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')

            # Run inference
            output_dict = sess.run(tensor_dict, feed_dict={image_tensor: np.expand_dims(image, 0)})

            # all outputs are float32 numpy arrays, so convert types as appropriate
            output_dict['filename'] = filename
            output_dict['num_detections'] = int(output_dict['num_detections'][0])
            output_dict['detection_classes'] = output_dict[
                'detection_classes'][0].astype(np.uint8)
            output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
            output_dict['detection_scores'] = output_dict['detection_scores'][0]
            if 'detection_masks' in output_dict:
                output_dict['detection_masks'] = output_dict['detection_masks'][0]
        return output_dict

def predict_image(TEST_IMAGE_PATHS, PATH_TO_CKPT, category_index, save_path):
    detection_graph = load_detection_graph(PATH_TO_CKPT)
    prediction_dict = defaultdict()
    start_time = time.time()
    for image_path in TEST_IMAGE_PATHS:
        toc = time.time()
        filename = image_path
        image = Image.open(image_path)
        # the array based representation of the image will be used later in order to prepare the
        # result image with boxes and labels on it.
        image_np = load_image_into_numpy_array(image)
        # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
        image_np_expanded = np.expand_dims(image_np, axis=0)
        # Actual detection.
        output_dict = run_inference_for_single_image(image_np, detection_graph, filename)
        # Visualization of the results of a detection.
        vis_util.visualize_boxes_and_labels_on_image_array(
            image_np,
            output_dict['detection_boxes'],
            output_dict['detection_classes'],
            output_dict['detection_scores'],
            category_index,
            instance_masks=output_dict.get('detection_masks'),
            use_normalized_coordinates=True,
            line_thickness=1)
        prediction_dict[filename] = output_dict
        plt.figure(figsize=(8,6), dpi=100)
        plt.imshow(image_np)
        plt.savefig(save_path+'{}'.format(filename))
        tic = time.time()
        print('{0} saved in {1:.2f}sec'.format(filename, tic-toc))
    end_time = time.time()
    print('{0:.2f}min to predict all images'.format((end_time-start_time)/60))
    with open('../predictions/predictions.pickle', 'wb') as f:
        pickle.dump(prediction_dict, f)
    return prediction_dict

Optimization Strategies (Ordered by Impact/Ease of Implementation) #

1. Remove Visualization & Saving Code Immediately

Your current code spends a significant chunk of time on vis_util.visualize_boxes_and_labels_on_image_array, plt.imshow, and plt.savefig—all of which you don't need. This is the quickest win to cut down per-image time.

Modified predict_image function:

def predict_image(TEST_IMAGE_PATHS, PATH_TO_CKPT):  # Remove unused params
    detection_graph = load_detection_graph(PATH_TO_CKPT)
    prediction_dict = defaultdict()
    start_time = time.time()
    for image_path in TEST_IMAGE_PATHS:
        toc = time.time()
        filename = image_path
        image = Image.open(image_path)
        image_np = load_image_into_numpy_array(image)
        # Actual detection.
        output_dict = run_inference_for_single_image(image_np, detection_graph, filename)
        prediction_dict[filename] = output_dict
        tic = time.time()
        print('{0} processed in {1:.2f}sec'.format(filename, tic-toc))
    end_time = time.time()
    print('{0:.2f}min to predict all images'.format((end_time-start_time)/60))
    with open('../predictions/predictions.pickle', 'wb') as f:
        pickle.dump(prediction_dict, f)
    return prediction_dict

2. Reuse the TensorFlow Session (Critical for Speed)

Right now, you're creating a new tf.Session() every time you run inference on a single image—this is extremely expensive (session initialization takes hundreds of milliseconds). Instead, create the session once and reuse it across all images.

Refactored code to reuse sessions:

def load_detection_graph_and_session(PATH_TO_CKPT):
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
        # Create session once here
        sess = tf.Session(graph=detection_graph)
    return detection_graph, sess

def run_inference_for_single_image(image, sess):  # Pass sess instead of graph
    # Cache tensor handles to avoid re-fetching every time
    tensor_dict = sess.graph.get_collection('inference_tensors')
    if not tensor_dict:
        ops = sess.graph.get_operations()
        all_tensor_names = {output.name for op in ops for output in op.outputs}
        tensor_dict = {}
        # Remove 'detection_masks' if you don't need them!
        for key in [
            'num_detections', 'detection_boxes', 'detection_scores',
            'detection_classes'
        ]:
            tensor_name = key + ':0'
            if tensor_name in all_tensor_names:
                tensor_dict[key] = sess.graph.get_tensor_by_name(tensor_name)
        # Cache for future calls
        sess.graph.add_to_collection('inference_tensors', tensor_dict)
    
    image_tensor = sess.graph.get_tensor_by_name('image_tensor:0')

    # Run inference
    output_dict = sess.run(tensor_dict, feed_dict={image_tensor: np.expand_dims(image, 0)})

    # Post-process outputs
    output_dict['num_detections'] = int(output_dict['num_detections'][0])
    output_dict['detection_classes'] = output_dict['detection_classes'][0].astype(np.uint8)
    output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
    output_dict['detection_scores'] = output_dict['detection_scores'][0]
    return output_dict

# Updated predict_image
def predict_image(TEST_IMAGE_PATHS, PATH_TO_CKPT):
    detection_graph, sess = load_detection_graph_and_session(PATH_TO_CKPT)
    prediction_dict = defaultdict()
    start_time = time.time()
    for image_path in TEST_IMAGE_PATHS:
        toc = time.time()
        filename = image_path
        image = Image.open(image_path)
        image_np = load_image_into_numpy_array(image)
        output_dict = run_inference_for_single_image(image_np, sess)
        output_dict['filename'] = filename
        prediction_dict[filename] = output_dict
        tic = time.time()
        print('{0} processed in {1:.2f}sec'.format(filename, tic-toc))
    # Close session when done
    sess.close()
    end_time = time.time()
    print('{0:.2f}min to predict all images'.format((end_time-start_time)/60))
    with open('../predictions/predictions.pickle', 'wb') as f:
        pickle.dump(prediction_dict, f)
    return prediction_dict

3. Eliminate Detection Mask Processing (If Unused)

Your code includes heavy logic for detection_masks—if you don't need masks, remove all references to them. Mask processing adds significant computation overhead, especially for single images. The refactored code above already excludes masks from the tensor list, which will save you time on slicing, reframing, and casting operations.

4. Batch Inference Instead of Single Image

GPUs are built for parallel processing. Feeding batches of images (e.g., 8, 16, or 32 at a time) instead of one will drastically improve throughput.

Batch inference implementation:

def run_inference_for_batch(images_batch, sess):
    # images_batch: numpy array of shape [batch_size, H, W, 3]
    tensor_dict = sess.graph.get_collection('inference_tensors')
    image_tensor = sess.graph.get_tensor_by_name('image_tensor:0')
    output_dict = sess.run(tensor_dict, feed_dict={image_tensor: images_batch})
    
    # Post-process batch outputs
    batch_output = []
    for i in range(len(images_batch)):
        single_output = {
            'num_detections': int(output_dict['num_detections'][i]),
            'detection_classes': output_dict['detection_classes'][i].astype(np.uint8),
            'detection_boxes': output_dict['detection_boxes'][i],
            'detection_scores': output_dict['detection_scores'][i]
        }
        batch_output.append(single_output)
    return batch_output

# Updated predict_image with batching
def predict_image(TEST_IMAGE_PATHS, PATH_TO_CKPT, batch_size=16):
    detection_graph, sess = load_detection_graph_and_session(PATH_TO_CKPT)
    prediction_dict = defaultdict()
    start_time = time.time()
    
    # Process images in batches
    for i in range(0, len(TEST_IMAGE_PATHS), batch_size):
        batch_paths = TEST_IMAGE_PATHS[i:i+batch_size]
        batch_images = []
        for image_path in batch_paths:
            image = Image.open(image_path)
            image_np = load_image_into_numpy_array(image)
            batch_images.append(image_np)
        batch_images_np = np.array(batch_images)
        
        # Run batch inference
        batch_outputs = run_inference_for_batch(batch_images_np, sess)
        
        # Map outputs to filenames
        for path, output in zip(batch_paths, batch_outputs):
            output['filename'] = path
            prediction_dict[path] = output
        print(f'Processed batch {i//batch_size +1} in {time.time()-toc:.2f}sec')
    
    sess.close()
    end_time = time.time()
    print('{0:.2f}min to predict all images'.format((end_time-start_time)/60))
    with open('../predictions/predictions.pickle', 'wb') as f:
        pickle.dump(prediction_dict, f)
    return prediction_dict

5. Optimize Image Loading/Preprocessing

PIL is slow for bulk image loading. Switch to OpenCV for faster image reading and conversion:

import cv2

def load_image_into_numpy_array(image_path):
    # Load directly with OpenCV (already in HWC format)
    image = cv2.imread(image_path)
    # Convert BGR to RGB (TF models expect RGB input)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image_rgb.astype(np.uint8)

This cuts down on image loading time significantly compared to PIL.

6. Use TensorRT for Model Optimization

Tesla K80 supports TensorRT, which optimizes TensorFlow models for GPU inference by fusing layers, quantizing weights, and optimizing memory usage. This can give you a 2-5x speedup on top of previous optimizations.

TensorRT optimization snippet:

from tensorflow.contrib.tensorrt.python import trt_convert as trt

def load_trt_optimized_graph(PATH_TO_CKPT, batch_size=16):
    # Load frozen graph
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def = tf.GraphDef()
        od_graph_def.ParseFromString(serialized_graph)
    
    # Convert to TensorRT optimized graph
    trt_graph = trt.create_inference_graph(
        input_graph_def=od_graph_def,
        outputs=['num_detections:0', 'detection_boxes:0', 'detection_scores:0', 'detection_classes:0'],
        max_batch_size=batch_size,
        max_workspace_size_bytes=1 << 30,  # 1GB workspace
        precision_mode='FP16'  # K80 supports FP16, faster than FP32
    )
    
    # Save optimized graph (optional)
    with tf.gfile.GFile('trt_optimized_graph.pb', 'wb') as f:
        f.write(trt_graph.SerializeToString())
    
    # Load and create session
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        tf.import_graph_def(trt_graph, name='')
        sess = tf.Session(graph=detection_graph)
    return detection_graph, sess

Use this function instead of load_detection_graph_and_session for massive speed gains.

7. Enable XLA Acceleration

XLA (Accelerated Linear Algebra) optimizes TensorFlow computation graphs for faster execution on GPUs. Enable it by modifying your session config:

config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
sess = tf.Session(graph=detection_graph, config=config)

This can add an extra 10-20% speedup on top of other optimizations.

Expected Speedup #

Combining all these optimizations:

• Removing visualization: ~2-3x speedup
• Reusing session: ~5-10x speedup
• Batch inference: ~2-3x speedup
• TensorRT optimization: ~2-5x speedup

Total speedup should easily exceed 20x, getting your per-image time down to ~0.09 seconds or better on the Tesla K80.

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