如何在NumPy或TensorFlow中高效实现批量图像裁剪

问题背景 #

适用场景为输入单张图像，需要按多组独立的起始x/y偏移量、裁剪高度、裁剪宽度参数批量生成子块。
朴素循环实现存在明显性能问题：

• NumPy下逐次切片、TensorFlow下循环调用裁剪/resize算子速度极慢，实测比C++实现慢20倍以上
• 尝试tf.vectorized_map、tf.while_loop等接口无明显提速
• 裁剪本质是内存拷贝操作，预分配内存场景下理论性能极高，性能瓶颈主要来自Python层循环调度、框架算子反复初始化的开销
  朴素实现代码如下（原代码存在切片变量笔误，已修正）：

import numpy as np
img = np.zeros([100, 100, 3])
ofsets_x = np.array([10, 15, 18])
img_w = np.array([10, 12, 15])
ofsets_y = np.array([20, 22, 14])
img_h = np.array([14, 12, 16])

crops= []
for i in range(ofsets_x.shape[0]):
    ofset_x = ofsets_x[i]
    ofset_y = ofsets_y[i]
    w = img_w[i]
    h = img_h[i]
    crop = img[ofset_x:ofset_x + w, ofset_y:ofset_y + h, :] 
    crops.append(crop)

NumPy 优化方案 #

核心思路是完全规避Python层循环，利用内存视图、批量索引一次性完成所有裁剪操作，性能接近原生memcpy：

• 若所有裁剪块尺寸一致，直接使用步长技巧生成滑动窗口视图，无额外内存开销，单次批量索引即可取出所有子块
• 若裁剪块尺寸不一致，先预分配输出数组内存，构造批量索引后一次性完成内存填充
  尺寸一致场景示例代码：

from numpy.lib.stride_tricks import sliding_window_view

def fast_batch_crop_numpy(img, ofsets_x, ofsets_y, crop_w, crop_h):
    H, W, C = img.shape
    n_crops = len(ofsets_x)
    # 生成全图滑动窗口视图，全程无数据拷贝
    window_view = sliding_window_view(img, (crop_h, crop_w), axis=(0, 1))
    # 单次高级索引完成所有裁剪，仅此处发生连续内存拷贝
    crops = window_view[ofsets_x, ofsets_y].transpose(0, 3, 4, 1, 2)[np.arange(n_crops), :, :, :, np.arange(C)]
    return crops

如果后续需要将所有子块resize到统一尺寸，不需要先裁出原始尺寸子块再做缩放，直接在滑动窗口上按采样率取点即可，可减少一次全量内存读写。

TensorFlow 优化方案 #

不要自行封装循环逻辑，直接调用框架内置的原生批量算子tf.image.crop_and_resize，该算子底层为C实现的批量内存拷贝+插值逻辑，内置SIMD优化，性能与手写C实现基本持平：

• 算子原生支持批量传入裁剪框坐标，一步完成裁剪+resize操作，无Python层调度开销
• 单图输入场景下仅需构造全0的框索引数组即可，不需要额外处理批次维度
  示例代码：

import tensorflow as tf

def fast_batch_crop_resize_tf(img, ofsets_x, ofsets_y, crop_w, crop_h, target_w, target_h):
    H, W, C = img.shape
    n_crops = len(ofsets_x)
    # 构造算子要求的归一化框坐标，格式为[y1, x1, y2, x2]，取值范围归一化到[0,1]
    boxes = tf.stack([
        ofsets_y / H,
        ofsets_x / W,
        (ofsets_y + crop_h) / H,
        (ofsets_x + crop_w) / W
    ], axis=1)
    # 所有裁剪块来自同一张输入图，框索引全部填0
    box_indices = tf.zeros(n_crops, dtype=tf.int32)
    # 单算子完成批量裁剪+缩放，无循环开销
    crops = tf.image.crop_and_resize(
        tf.expand_dims(img, 0),
        boxes,
        box_indices,
        crop_size=(target_h, target_w),
        method='bilinear'
    )
    return crops

如果不需要resize仅需原始尺寸裁剪结果，可使用tf.gather_nd构造批量坐标索引一次性取数，性能远高于循环调用tf.slice。

性能优化注意事项 #

• Python层for循环的调度开销远大于裁剪操作本身，批量裁剪场景下要100%规避逐次循环调用单裁剪算子的写法
• 裁剪+resize的组合逻辑不要拆分为两个独立步骤，合并为单个算子调用可减少一次全量内存读写，性能提升可达100%
• 优先使用框架内置的批量原生算子，这类算子均做过内存预分配、缓存对齐、SIMD指令优化，性能不会低于手写C++实现

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